dictionary.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
 
 
#pragma once
 
// ---------------------------------------------------------------------------
// dictionary.hpp -- Dictionary encoding for Parquet
//
// Implements PLAIN_DICTIONARY (encoding=2) for dictionary pages and
// RLE_DICTIONARY (encoding=8) for data pages. Critical for low-cardinality
// string columns (symbols, sides, exchanges) where 10-50x compression is
// typical.
//
// Wire format:
//   Dictionary page: PLAIN-encoded unique values
//     - BYTE_ARRAY: 4-byte LE length + bytes per entry
//     - INT32/FLOAT: 4-byte LE per entry
//     - INT64/DOUBLE: 8-byte LE per entry
//
//   Data page: 1-byte bit_width + RLE/Bit-Pack hybrid encoded indices
//     - bit_width = 0 for dict_size==1, else ceil(log2(dict_size))
//     - Indices encoded using RLE/Bit-Packing Hybrid (see rle.hpp)
// ---------------------------------------------------------------------------
 
#include "signet/encoding/rle.hpp"
#include "signet/error.hpp"
#include "signet/types.hpp"
 
#include <cassert>
#include <cstdint>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace signet::forge {
 
// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
 
namespace detail {
 
inline int dict_bit_width(size_t dict_size) {
    if (dict_size <= 1) return 0;
    // ceil(log2(n)) = number of bits needed to represent 0..n-1
    int bits = 0;
    size_t n = dict_size - 1; // max index value
    while (n > 0) {
        ++bits;
        n >>= 1;
    }
    return bits;
}
 
// -- PLAIN encoding helpers for dictionary page entries ----------------------
 
inline void plain_encode_value(std::vector<uint8_t>& buf, const std::string& val) {
    if (val.size() > UINT32_MAX) throw std::length_error("String exceeds Parquet BYTE_ARRAY 4 GB limit");
    uint32_t len = static_cast<uint32_t>(val.size());
    buf.push_back(static_cast<uint8_t>((len      ) & 0xFF));
    buf.push_back(static_cast<uint8_t>((len >>  8) & 0xFF));
    buf.push_back(static_cast<uint8_t>((len >> 16) & 0xFF));
    buf.push_back(static_cast<uint8_t>((len >> 24) & 0xFF));
    buf.insert(buf.end(),
               reinterpret_cast<const uint8_t*>(val.data()),
               reinterpret_cast<const uint8_t*>(val.data()) + val.size());
}
 
inline void plain_encode_value(std::vector<uint8_t>& buf, int32_t val) {
    uint32_t bits;
    std::memcpy(&bits, &val, sizeof(bits));
    buf.push_back(static_cast<uint8_t>((bits      ) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >>  8) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >> 16) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >> 24) & 0xFF));
}
 
inline void plain_encode_value(std::vector<uint8_t>& buf, int64_t val) {
    uint64_t bits;
    std::memcpy(&bits, &val, sizeof(bits));
    for (int i = 0; i < 8; ++i) {
        buf.push_back(static_cast<uint8_t>(bits & 0xFF));
        bits >>= 8;
    }
}
 
inline void plain_encode_value(std::vector<uint8_t>& buf, float val) {
    uint32_t bits;
    std::memcpy(&bits, &val, sizeof(bits));
    buf.push_back(static_cast<uint8_t>((bits      ) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >>  8) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >> 16) & 0xFF));
    buf.push_back(static_cast<uint8_t>((bits >> 24) & 0xFF));
}
 
inline void plain_encode_value(std::vector<uint8_t>& buf, double val) {
    uint64_t bits;
    std::memcpy(&bits, &val, sizeof(bits));
    for (int i = 0; i < 8; ++i) {
        buf.push_back(static_cast<uint8_t>(bits & 0xFF));
        bits >>= 8;
    }
}
 
// -- PLAIN decoding helpers for dictionary page entries ----------------------
 
inline std::string plain_decode_value(const uint8_t* data, size_t& pos,
                                       size_t size, std::string* /*tag*/) {
    if (pos + 4 > size) return {};
    uint32_t len;
    std::memcpy(&len, data + pos, 4);
    pos += 4;
    if (pos + len > size) return {};
    std::string val(reinterpret_cast<const char*>(data + pos), len);
    pos += len;
    return val;
}
 
inline int32_t plain_decode_value(const uint8_t* data, size_t& pos,
                                   size_t size, int32_t* /*tag*/) {
    if (pos + 4 > size) return 0;
    int32_t val;
    std::memcpy(&val, data + pos, 4);
    pos += 4;
    return val;
}
 
inline int64_t plain_decode_value(const uint8_t* data, size_t& pos,
                                   size_t size, int64_t* /*tag*/) {
    if (pos + 8 > size) return 0;
    int64_t val;
    std::memcpy(&val, data + pos, 8);
    pos += 8;
    return val;
}
 
inline float plain_decode_value(const uint8_t* data, size_t& pos,
                                 size_t size, float* /*tag*/) {
    if (pos + 4 > size) return 0.0f;
    float val;
    std::memcpy(&val, data + pos, 4);
    pos += 4;
    return val;
}
 
inline double plain_decode_value(const uint8_t* data, size_t& pos,
                                  size_t size, double* /*tag*/) {
    if (pos + 8 > size) return 0.0;
    double val;
    std::memcpy(&val, data + pos, 8);
    pos += 8;
    return val;
}
 
} // namespace detail
 
// ===========================================================================
// DictionaryEncoder
// ===========================================================================
//
// Builds a dictionary of unique values and encodes data as indices into
// that dictionary. The dictionary page is PLAIN-encoded, and the indices
// page uses RLE_DICTIONARY (bit_width byte + RLE/Bit-Pack hybrid).
//
// Usage:
//     DictionaryEncoder<std::string> enc;
//     for (auto& s : strings) enc.put(s);
//     enc.flush();
//     auto dict_page = enc.dictionary_page();   // PLAIN-encoded unique values
//     auto idx_page  = enc.indices_page();       // bit_width + RLE-encoded indices
//
template <typename T>
class DictionaryEncoder {
public:
    DictionaryEncoder() = default;
 
    static constexpr size_t MAX_DICTIONARY_ENTRIES = 1 << 20; // 1M entries
 
    [[nodiscard]] bool is_full() const { return dict_values_.size() >= MAX_DICTIONARY_ENTRIES; }
 
    bool put(const T& value) {
        auto it = dict_map_.find(value);
        uint32_t index;
        if (it == dict_map_.end()) {
            if (dict_values_.size() >= MAX_DICTIONARY_ENTRIES) {
                return false; // dictionary full — caller should fall back to PLAIN
            }
            index = static_cast<uint32_t>(dict_values_.size());
            dict_map_.emplace(value, index);
            dict_values_.push_back(value);
        } else {
            index = it->second;
        }
        indices_.push_back(index);
        return true;
    }
 
    void flush() {
        // Nothing to accumulate -- indices are stored incrementally.
        // This exists for API symmetry with other encoders.
    }
 
    [[nodiscard]] std::vector<uint8_t> dictionary_page() const {
        std::vector<uint8_t> buf;
        for (const auto& val : dict_values_) {
            detail::plain_encode_value(buf, val);
        }
        return buf;
    }
 
    [[nodiscard]] std::vector<uint8_t> indices_page() const {
        int bw = bit_width();
 
        // RLE-encode the indices
        auto rle_payload = RleEncoder::encode(indices_.data(),
                                               indices_.size(), bw);
 
        // Prepend the bit_width byte
        std::vector<uint8_t> result;
        result.reserve(1 + rle_payload.size());
        result.push_back(static_cast<uint8_t>(bw));
        result.insert(result.end(), rle_payload.begin(), rle_payload.end());
        return result;
    }
 
    [[nodiscard]] size_t dictionary_size() const { return dict_values_.size(); }
 
    [[nodiscard]] size_t num_values() const { return indices_.size(); }
 
    [[nodiscard]] int bit_width() const {
        return detail::dict_bit_width(dict_values_.size());
    }
 
    void reset() {
        dict_map_.clear();
        dict_values_.clear();
        indices_.clear();
    }
 
    [[nodiscard]] bool is_worthwhile() const {
        if (indices_.empty()) return false;
        return dict_values_.size() < indices_.size() * 4 / 10;
    }
 
private:
    std::unordered_map<T, uint32_t> dict_map_;   
    std::vector<T>                  dict_values_; 
    std::vector<uint32_t>           indices_;     
};
 
// ===========================================================================
// DictionaryDecoder
// ===========================================================================
//
// Decodes a dictionary page (PLAIN-encoded values) and an indices page
// (bit_width byte + RLE-encoded indices) back to the original typed values.
//
// Usage:
//     DictionaryDecoder<std::string> dec(dict_data, dict_size, num_entries, type);
//     auto values = dec.decode(idx_data, idx_size, num_values);
//
template <typename T>
class DictionaryDecoder {
public:
    DictionaryDecoder(const uint8_t* dict_data, size_t dict_size,
                      size_t num_dict_entries, PhysicalType type)
        : type_(type)
    {
        // Decode dictionary entries from PLAIN-encoded data
        dict_values_.reserve((std::min)(static_cast<size_t>(num_dict_entries), dict_size));
        size_t pos = 0;
        for (size_t i = 0; i < num_dict_entries; ++i) {
            T* tag = nullptr;
            dict_values_.push_back(
                detail::plain_decode_value(dict_data, pos, dict_size, tag));
        }
    }
 
    [[nodiscard]] expected<std::vector<T>> decode(const uint8_t* indices_data,
                                                   size_t indices_size,
                                                   size_t num_values) const {
        if (indices_size == 0 || num_values == 0) return std::vector<T>{};
 
        // First byte is the bit_width
        int bw = static_cast<int>(indices_data[0]);
 
        // Remaining bytes are the RLE-encoded indices
        auto indices = RleDecoder::decode(indices_data + 1,
                                           indices_size - 1,
                                           bw, num_values);
 
        // Map indices back to dictionary values
        std::vector<T> result;
        result.reserve(indices.size());
        for (uint32_t idx : indices) {
            if (static_cast<size_t>(idx) >= dict_values_.size()) {
                return Error{ErrorCode::CORRUPT_DATA,
                             "dictionary index " + std::to_string(idx)
                             + " out of range (dict size="
                             + std::to_string(dict_values_.size()) + ")"};
            }
            result.push_back(dict_values_[idx]);
        }
        return result;
    }
 
    [[nodiscard]] size_t dictionary_size() const { return dict_values_.size(); }
 
private:
    PhysicalType    type_;         
    std::vector<T>  dict_values_;  
};
 
} // namespace signet::forge