column_index.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include "signet/thrift/compact.hpp"
#include "signet/statistics.hpp"    // from_le_bytes<T>
 
#include <cstdint>
#include <cstring>
#include <string>
#include <vector>
 
namespace signet::forge {
 
struct PageLocation {
    int64_t offset              = 0;  
    int32_t compressed_page_size = 0; 
    int64_t first_row_index     = 0;  
 
    void serialize(thrift::CompactEncoder& enc) const {
        enc.begin_struct();
 
        // field 1: offset (i64)
        enc.write_field(1, thrift::compact_type::I64);
        enc.write_i64(offset);
 
        // field 2: compressed_page_size (i32)
        enc.write_field(2, thrift::compact_type::I32);
        enc.write_i32(compressed_page_size);
 
        // field 3: first_row_index (i64)
        enc.write_field(3, thrift::compact_type::I64);
        enc.write_i64(first_row_index);
 
        enc.write_stop();
        enc.end_struct();
    }
 
    void deserialize(thrift::CompactDecoder& dec) {
        dec.begin_struct();
        for (;;) {
            auto [fid, ftype] = dec.read_field_header();
            if (ftype == thrift::compact_type::STOP) break;
            switch (fid) {
                case 1: offset               = dec.read_i64(); break;
                case 2: compressed_page_size  = dec.read_i32(); break;
                case 3: first_row_index      = dec.read_i64(); break;
                default: dec.skip_field(ftype); break;
            }
        }
        dec.end_struct();
    }
};
 
struct OffsetIndex {
    bool valid_ = true; 
    std::vector<PageLocation> page_locations; 
 
    [[nodiscard]] bool valid() const { return valid_; }
 
    void serialize(thrift::CompactEncoder& enc) const {
        enc.begin_struct();
 
        // field 1: page_locations (list<struct>)
        enc.write_field(1, thrift::compact_type::LIST);
        enc.write_list_header(thrift::compact_type::STRUCT,
                              static_cast<int32_t>(page_locations.size()));
        for (const auto& loc : page_locations) {
            loc.serialize(enc);
        }
 
        enc.write_stop();
        enc.end_struct();
    }
 
    void deserialize(thrift::CompactDecoder& dec) {
        dec.begin_struct();
        for (;;) {
            auto [fid, ftype] = dec.read_field_header();
            if (ftype == thrift::compact_type::STOP) break;
            switch (fid) {
                case 1: {
                    auto [elem_type, count] = dec.read_list_header();
                    // CWE-400: Uncontrolled Resource Consumption — 10M cap on list counts
                    if (count < 0 || static_cast<size_t>(count) > 10'000'000) {
                        valid_ = false;
                        return; // list count exceeds 10M cap or is negative
                    }
                    page_locations.resize(static_cast<size_t>(count));
                    for (int32_t i = 0; i < count; ++i) {
                        page_locations[static_cast<size_t>(i)].deserialize(dec);
                    }
                    break;
                }
                default: dec.skip_field(ftype); break;
            }
        }
        dec.end_struct();
    }
};
 
struct ColumnIndex {
    bool                     valid_ = true; 
    std::vector<bool>        null_pages;  
    std::vector<std::string> min_values;  
    std::vector<std::string> max_values;  
 
    enum class BoundaryOrder : int32_t {
        UNORDERED  = 0, 
        ASCENDING  = 1, 
        DESCENDING = 2  
    };
    BoundaryOrder boundary_order = BoundaryOrder::UNORDERED; 
 
    std::vector<int64_t> null_counts; 
 
    [[nodiscard]] bool valid() const { return valid_; }
 
    void serialize(thrift::CompactEncoder& enc) const {
        enc.begin_struct();
 
        // field 1: null_pages (list<bool>)
        enc.write_field(1, thrift::compact_type::LIST);
        enc.write_list_header(thrift::compact_type::BOOL_TRUE,
                              static_cast<int32_t>(null_pages.size()));
        for (bool np : null_pages) {
            enc.write_bool(np);
        }
 
        // field 2: min_values (list<binary>)
        enc.write_field(2, thrift::compact_type::LIST);
        enc.write_list_header(thrift::compact_type::BINARY,
                              static_cast<int32_t>(min_values.size()));
        for (const auto& mv : min_values) {
            enc.write_string(mv);
        }
 
        // field 3: max_values (list<binary>)
        enc.write_field(3, thrift::compact_type::LIST);
        enc.write_list_header(thrift::compact_type::BINARY,
                              static_cast<int32_t>(max_values.size()));
        for (const auto& mv : max_values) {
            enc.write_string(mv);
        }
 
        // field 4: boundary_order (i32)
        enc.write_field(4, thrift::compact_type::I32);
        enc.write_i32(static_cast<int32_t>(boundary_order));
 
        // field 5: null_counts (list<i64>, optional -- written only if non-empty)
        if (!null_counts.empty()) {
            enc.write_field(5, thrift::compact_type::LIST);
            enc.write_list_header(thrift::compact_type::I64,
                                  static_cast<int32_t>(null_counts.size()));
            for (int64_t nc : null_counts) {
                enc.write_i64(nc);
            }
        }
 
        enc.write_stop();
        enc.end_struct();
    }
 
    void deserialize(thrift::CompactDecoder& dec) {
        dec.begin_struct();
        for (;;) {
            auto [fid, ftype] = dec.read_field_header();
            if (ftype == thrift::compact_type::STOP) break;
            switch (fid) {
                case 1: {
                    // null_pages: list<bool>
                    auto [elem_type, count] = dec.read_list_header();
                    // CWE-400: Uncontrolled Resource Consumption — 10M cap on list counts
                    if (count < 0 || static_cast<size_t>(count) > 10'000'000) {
                        valid_ = false; return; // list count exceeds 10M cap or is negative
                    }
                    null_pages.resize(static_cast<size_t>(count));
                    for (int32_t i = 0; i < count; ++i) {
                        null_pages[static_cast<size_t>(i)] = dec.read_bool();
                    }
                    break;
                }
                case 2: {
                    // min_values: list<binary>
                    auto [elem_type, count] = dec.read_list_header();
                    // CWE-400: Uncontrolled Resource Consumption — 10M cap on list counts
                    if (count < 0 || static_cast<size_t>(count) > 10'000'000) {
                        valid_ = false; return; // list count exceeds 10M cap or is negative
                    }
                    min_values.resize(static_cast<size_t>(count));
                    for (int32_t i = 0; i < count; ++i) {
                        min_values[static_cast<size_t>(i)] = dec.read_string();
                    }
                    break;
                }
                case 3: {
                    // max_values: list<binary>
                    auto [elem_type, count] = dec.read_list_header();
                    // CWE-400: Uncontrolled Resource Consumption — 10M cap on list counts
                    if (count < 0 || static_cast<size_t>(count) > 10'000'000) {
                        valid_ = false; return; // list count exceeds 10M cap or is negative
                    }
                    max_values.resize(static_cast<size_t>(count));
                    for (int32_t i = 0; i < count; ++i) {
                        max_values[static_cast<size_t>(i)] = dec.read_string();
                    }
                    break;
                }
                case 4:
                    boundary_order = static_cast<BoundaryOrder>(dec.read_i32());
                    break;
                case 5: {
                    // null_counts: list<i64>
                    auto [elem_type, count] = dec.read_list_header();
                    // CWE-400: Uncontrolled Resource Consumption — 10M cap on list counts
                    if (count < 0 || static_cast<size_t>(count) > 10'000'000) {
                        valid_ = false; return; // list count exceeds 10M cap or is negative
                    }
                    null_counts.resize(static_cast<size_t>(count));
                    for (int32_t i = 0; i < count; ++i) {
                        null_counts[static_cast<size_t>(i)] = dec.read_i64();
                    }
                    break;
                }
                default: dec.skip_field(ftype); break;
            }
        }
        dec.end_struct();
    }
 
    [[nodiscard]] std::vector<size_t> filter_pages(
            const std::string& min_val,
            const std::string& max_val,
            PhysicalType physical_type = PhysicalType::BYTE_ARRAY) const {
 
        // Typed comparison helper: returns negative/zero/positive like strcmp.
        auto typed_compare = [&](const std::string& a, const std::string& b) -> int {
            switch (physical_type) {
                case PhysicalType::INT32: {
                    if (a.size() >= sizeof(int32_t) && b.size() >= sizeof(int32_t)) {
                        auto va = from_le_bytes<int32_t>({reinterpret_cast<const uint8_t*>(a.data()),
                                                          reinterpret_cast<const uint8_t*>(a.data()) + a.size()});
                        auto vb = from_le_bytes<int32_t>({reinterpret_cast<const uint8_t*>(b.data()),
                                                          reinterpret_cast<const uint8_t*>(b.data()) + b.size()});
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                }
                case PhysicalType::INT64: {
                    if (a.size() >= sizeof(int64_t) && b.size() >= sizeof(int64_t)) {
                        auto va = from_le_bytes<int64_t>({reinterpret_cast<const uint8_t*>(a.data()),
                                                          reinterpret_cast<const uint8_t*>(a.data()) + a.size()});
                        auto vb = from_le_bytes<int64_t>({reinterpret_cast<const uint8_t*>(b.data()),
                                                          reinterpret_cast<const uint8_t*>(b.data()) + b.size()});
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                }
                case PhysicalType::FLOAT: {
                    if (a.size() >= sizeof(float) && b.size() >= sizeof(float)) {
                        auto va = from_le_bytes<float>({reinterpret_cast<const uint8_t*>(a.data()),
                                                        reinterpret_cast<const uint8_t*>(a.data()) + a.size()});
                        auto vb = from_le_bytes<float>({reinterpret_cast<const uint8_t*>(b.data()),
                                                        reinterpret_cast<const uint8_t*>(b.data()) + b.size()});
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                }
                case PhysicalType::DOUBLE: {
                    if (a.size() >= sizeof(double) && b.size() >= sizeof(double)) {
                        auto va = from_le_bytes<double>({reinterpret_cast<const uint8_t*>(a.data()),
                                                         reinterpret_cast<const uint8_t*>(a.data()) + a.size()});
                        auto vb = from_le_bytes<double>({reinterpret_cast<const uint8_t*>(b.data()),
                                                         reinterpret_cast<const uint8_t*>(b.data()) + b.size()});
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                }
                default:
                    break;
            }
            // Fallback: lexicographic comparison
            return (a < b) ? -1 : (a > b) ? 1 : 0;
        };
 
        std::vector<size_t> matching;
        size_t num_pages = min_values.size();
 
        for (size_t i = 0; i < num_pages; ++i) {
            // Skip all-null pages -- they cannot contain any matching data
            if (i < null_pages.size() && null_pages[i]) {
                continue;
            }
 
            // Skip if page max < query min (entire page below the range)
            if (i < max_values.size() && typed_compare(max_values[i], min_val) < 0) {
                continue;
            }
 
            // Skip if page min > query max (entire page above the range)
            if (i < min_values.size() && typed_compare(min_values[i], max_val) > 0) {
                continue;
            }
 
            matching.push_back(i);
        }
 
        return matching;
    }
};
 
class ColumnIndexBuilder {
public:
    void start_page() {
        pages_.emplace_back();
    }
 
    void set_min(const std::string& min_val) {
        if (!pages_.empty()) {
            pages_.back().min_value = min_val;
        }
    }
 
    void set_max(const std::string& max_val) {
        if (!pages_.empty()) {
            pages_.back().max_value = max_val;
        }
    }
 
    void set_null_page(bool is_null) {
        if (!pages_.empty()) {
            pages_.back().null_page = is_null;
        }
    }
 
    void set_null_count(int64_t count) {
        if (!pages_.empty()) {
            pages_.back().null_count = count;
        }
    }
 
    void set_first_row_index(int64_t row_index) {
        if (!pages_.empty()) {
            pages_.back().first_row_index = row_index;
        }
    }
 
    void set_page_location(int64_t offset, int32_t compressed_size) {
        if (!pages_.empty()) {
            pages_.back().offset = offset;
            pages_.back().compressed_size = compressed_size;
        }
    }
 
    [[nodiscard]] ColumnIndex build_column_index(
            PhysicalType pt = PhysicalType::BYTE_ARRAY) const {
        ColumnIndex ci;
        ci.null_pages.reserve(pages_.size());
        ci.min_values.reserve(pages_.size());
        ci.max_values.reserve(pages_.size());
        ci.null_counts.reserve(pages_.size());
 
        bool has_any_null_counts = false;
        for (const auto& p : pages_) {
            ci.null_pages.push_back(p.null_page);
            ci.min_values.push_back(p.min_value);
            ci.max_values.push_back(p.max_value);
            ci.null_counts.push_back(p.null_count);
            if (p.null_count > 0) {
                has_any_null_counts = true;
            }
        }
 
        // Determine boundary order by scanning min_values (type-aware)
        ci.boundary_order = detect_boundary_order(ci.min_values, pt);
 
        // Only include null_counts if at least one page has nulls,
        // or if the caller explicitly set null counts. The Parquet spec
        // makes null_counts optional, so we always include them for
        // completeness (readers expect them for null-page detection).
        (void)has_any_null_counts;
 
        return ci;
    }
 
    [[nodiscard]] OffsetIndex build_offset_index() const {
        OffsetIndex oi;
        oi.page_locations.reserve(pages_.size());
 
        for (const auto& p : pages_) {
            PageLocation loc;
            loc.offset               = p.offset;
            loc.compressed_page_size = p.compressed_size;
            loc.first_row_index      = p.first_row_index;
            oi.page_locations.push_back(loc);
        }
 
        return oi;
    }
 
    void reset() {
        pages_.clear();
    }
 
    [[nodiscard]] size_t num_pages() const { return pages_.size(); }
 
private:
    struct PageInfo {
        std::string min_value;
        std::string max_value;
        bool        null_page       = false;
        int64_t     null_count      = 0;
        int64_t     first_row_index = 0;
        int64_t     offset          = 0;
        int32_t     compressed_size = 0;
    };
 
    std::vector<PageInfo> pages_;
 
    [[nodiscard]] static ColumnIndex::BoundaryOrder detect_boundary_order(
            const std::vector<std::string>& values,
            PhysicalType pt = PhysicalType::BYTE_ARRAY) {
 
        if (values.size() <= 1) {
            return ColumnIndex::BoundaryOrder::ASCENDING;
        }
 
        // Type-aware comparator: returns <0, 0, >0 like strcmp
        auto typed_cmp = [pt](const std::string& a, const std::string& b) -> int {
            switch (pt) {
                case PhysicalType::INT32:
                    if (a.size() >= sizeof(int32_t) && b.size() >= sizeof(int32_t)) {
                        int32_t va, vb;
                        std::memcpy(&va, a.data(), sizeof(int32_t));
                        std::memcpy(&vb, b.data(), sizeof(int32_t));
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                case PhysicalType::INT64:
                    if (a.size() >= sizeof(int64_t) && b.size() >= sizeof(int64_t)) {
                        int64_t va, vb;
                        std::memcpy(&va, a.data(), sizeof(int64_t));
                        std::memcpy(&vb, b.data(), sizeof(int64_t));
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                case PhysicalType::FLOAT:
                    if (a.size() >= sizeof(float) && b.size() >= sizeof(float)) {
                        float va, vb;
                        std::memcpy(&va, a.data(), sizeof(float));
                        std::memcpy(&vb, b.data(), sizeof(float));
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                case PhysicalType::DOUBLE:
                    if (a.size() >= sizeof(double) && b.size() >= sizeof(double)) {
                        double va, vb;
                        std::memcpy(&va, a.data(), sizeof(double));
                        std::memcpy(&vb, b.data(), sizeof(double));
                        return (va < vb) ? -1 : (va > vb) ? 1 : 0;
                    }
                    break;
                default:
                    break;
            }
            // Fallback: lexicographic for BYTE_ARRAY, BOOLEAN, etc.
            return a.compare(b);
        };
 
        bool ascending  = true;
        bool descending = true;
 
        for (size_t i = 1; i < values.size(); ++i) {
            int cmp = typed_cmp(values[i], values[i - 1]);
            if (cmp < 0) ascending  = false;
            if (cmp > 0) descending = false;
            if (!ascending && !descending) break;
        }
 
        if (ascending) return ColumnIndex::BoundaryOrder::ASCENDING;
        if (descending) return ColumnIndex::BoundaryOrder::DESCENDING;
        return ColumnIndex::BoundaryOrder::UNORDERED;
    }
};
 
} // namespace signet::forge