feat: support fixed length chunked dictionary for rangebitmap (#167)

Author: fafacao86

src/paimon/CMakeLists.txt

@@ -366,6 +366,7 @@ if(PAIMON_BUILD_TESTS)
                     common/file_index/bitmap/apply_bitmap_index_batch_reader_test.cpp
                     common/file_index/bsi/bit_slice_index_bitmap_file_index_test.cpp
                     common/file_index/bsi/bit_slice_index_roaring_bitmap_test.cpp
+                    common/file_index/rangebitmap/dictionary/chunked_dictionary_test.cpp
                     common/file_index/bloomfilter/bloom_filter_file_index_test.cpp
                     common/file_index/bloomfilter/fast_hash_test.cpp
                     common/global_index/complete_index_score_batch_reader_test.cpp

src/paimon/common/file_index/CMakeLists.txt

@@ -23,7 +23,11 @@ set(PAIMON_FILE_INDEX_SRC
     bsi/bit_slice_index_roaring_bitmap.cpp
     bloomfilter/bloom_filter_file_index.cpp
     bloomfilter/bloom_filter_file_index_factory.cpp
-    bloomfilter/fast_hash.cpp)
+    bloomfilter/fast_hash.cpp
+    rangebitmap/dictionary/chunked_dictionary.cpp
+    rangebitmap/dictionary/fixed_length_chunk.cpp
+    rangebitmap/dictionary/key_factory.cpp
+    rangebitmap/utils/literal_serialization_utils.cpp)
 
 add_paimon_lib(paimon_file_index
                SOURCES

src/paimon/common/file_index/rangebitmap/dictionary/chunk.h

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2026-present Alibaba Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <fmt/format.h>
+
+#include <memory>
+
+#include "paimon/memory/bytes.h"
+#include "paimon/predicate/literal.h"
+#include "paimon/result.h"
+
+namespace paimon {
+
+class Chunk {
+ public:
+    virtual ~Chunk() = default;
+
+    virtual Result<bool> TryAdd(const Literal& key) = 0;
+
+    /// Finds a key within this chunk using binary search.
+    ///
+    /// @param key The key to search for
+    /// @return index of the search key, if it is contained in the chunk;
+    ///         otherwise, -(insertion_point + 1). The insertion point
+    ///         is defined as the point at which the key would be inserted into the chunk:
+    ///         the index of the first element greater than the key, or the chunk's size
+    ///         if all elements in the chunk are less than the specified key. Note that
+    ///         this guarantees that the return value will be >= 0 if and only if the key is found.
+    virtual Result<int32_t> Find(const Literal& key) {
+        PAIMON_ASSIGN_OR_RAISE(int32_t cmp_with_key, CompareKey(Key(), key));
+        if (cmp_with_key == 0) {
+            return Code();
+        }
+        int32_t low = 0;
+        int32_t high = Size() - 1;
+        const int32_t base = Code() + 1;
+        while (low <= high) {
+            const int32_t mid = low + (high - low) / 2;
+            PAIMON_ASSIGN_OR_RAISE(Literal key_at_mid, GetKey(mid));
+            PAIMON_ASSIGN_OR_RAISE(int32_t cmp, CompareKey(key_at_mid, key));
+            if (cmp < 0) {
+                low = mid + 1;
+            } else if (cmp > 0) {
+                high = mid - 1;
+            } else {
+                return base + mid;
+            }
+        }
+        return -(base + low + 1);
+    }
+
+    virtual Result<Literal> Find(int32_t code) {
+        const auto current = Code();
+        if (current == code) {
+            return Key();
+        }
+        const auto index = code - current - 1;
+        if (index < 0 || index >= Size()) {
+            return Status::Invalid(fmt::format("Invalid Code: {}", code));
+        }
+        return GetKey(index);
+    }
+
+    virtual const Literal& Key() const = 0;
+
+    virtual int32_t Code() const = 0;
+
+    virtual int32_t Offset() const = 0;
+
+    virtual void SetOffset(int32_t offset) = 0;
+
+    virtual int32_t Size() const = 0;
+
+    virtual Result<PAIMON_UNIQUE_PTR<Bytes>> SerializeChunk() const = 0;
+
+    virtual Result<PAIMON_UNIQUE_PTR<Bytes>> SerializeKeys() const = 0;
+
+ protected:
+    virtual Result<Literal> GetKey(int32_t index) = 0;
+
+    virtual Result<int32_t> CompareKey(const Literal& lhs, const Literal& rhs) = 0;
+};
+
+}  // namespace paimon

src/paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.cpp

@@ -0,0 +1,247 @@
+/*
+ * Copyright 2026-present Alibaba Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/common/file_index/rangebitmap/dictionary/chunked_dictionary.h"
+
+#include <algorithm>
+
+#include "fmt/format.h"
+#include "paimon/common/file_index/rangebitmap/dictionary/fixed_length_chunk.h"
+#include "paimon/common/file_index/rangebitmap/dictionary/key_factory.h"
+#include "paimon/common/memory/memory_segment_utils.h"
+#include "paimon/fs/file_system.h"
+#include "paimon/io/byte_array_input_stream.h"
+#include "paimon/io/data_input_stream.h"
+#include "paimon/memory/bytes.h"
+#include "paimon/result.h"
+#include "paimon/status.h"
+
+namespace paimon {
+
+Result<int32_t> ChunkedDictionary::Find(const Literal& key) {
+    int32_t low = 0;
+    int32_t high = size_ - 1;
+    while (low <= high) {
+        const int32_t mid = low + (high - low) / 2;
+        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Chunk> chunk, GetChunk(mid));
+        PAIMON_ASSIGN_OR_RAISE(const int32_t result, factory_->CompareLiteral(chunk->Key(), key));
+        if (result > 0) {
+            high = mid - 1;
+        } else if (result < 0) {
+            low = mid + 1;
+        } else {
+            return chunk->Code();
+        }
+    }
+    if (low == 0) {
+        return -(low + 1);  // This makes sure the return value is negative if not found
+    }
+    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Chunk> prev_chunk, GetChunk(low - 1));
+    return prev_chunk->Find(key);
+}
+
+Result<Literal> ChunkedDictionary::Find(int32_t code) {
+    if (size_ <= 0) {
+        return Status::Invalid(fmt::format("Cannot find code {} in an empty Dictionary", code));
+    }
+    if (code < 0) {
+        return Status::Invalid(fmt::format("Invalid code: {}", code));
+    }
+    int32_t low = 0;
+    int32_t high = size_ - 1;
+
+    while (low <= high) {
+        const int32_t mid = low + (high - low) / 2;
+        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Chunk> chunk, GetChunk(mid));
+
+        int32_t chunk_code = chunk->Code();
+        if (chunk_code > code) {
+            high = mid - 1;
+        } else if (chunk_code < code) {
+            low = mid + 1;
+        } else {
+            return {chunk->Key()};
+        }
+    }
+    if (low == 0) {
+        return Status::Invalid(fmt::format("Cannot find code {}", code));
+    }
+    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Chunk> prev_chunk, GetChunk(low - 1));
+    return prev_chunk->Find(code);
+}
+
+Result<std::shared_ptr<Chunk>> ChunkedDictionary::GetChunk(int32_t index) {
+    if (index < 0 || index >= size_) {
+        return Status::Invalid(fmt::format("Invalid chunk index: {}", index));
+    }
+    if (offsets_bytes_ == nullptr || chunks_bytes_ == nullptr) {
+        PAIMON_RETURN_NOT_OK(input_stream_->Seek(body_offset_, FS_SEEK_SET));
+        auto offsets = Bytes::AllocateBytes(offsets_length_, pool_.get());
+        PAIMON_RETURN_NOT_OK(input_stream_->Read(offsets->data(), offsets_length_));
+        offsets_bytes_ = std::move(offsets);
+        auto chunks = Bytes::AllocateBytes(chunks_length_, pool_.get());
+        PAIMON_RETURN_NOT_OK(input_stream_->Read(chunks->data(), chunks_length_));
+        chunks_bytes_ = std::move(chunks);
+    }
+    if (chunks_cache_[index]) {
+        return chunks_cache_[index];
+    }
+    auto data_in = std::make_unique<DataInputStream>(
+        std::make_shared<ByteArrayInputStream>(offsets_bytes_->data(), offsets_length_));
+    PAIMON_RETURN_NOT_OK(data_in->Seek(sizeof(int32_t) * index));
+    PAIMON_ASSIGN_OR_RAISE(int32_t chunk_offset, data_in->ReadValue<int32_t>());
+    PAIMON_ASSIGN_OR_RAISE(
+        std::unique_ptr<Chunk> chunk,
+        factory_->MmapChunk(input_stream_, body_offset_ + offsets_length_ + chunk_offset,
+                            body_offset_ + chunks_length_ + offsets_length_, pool_));
+    chunks_cache_[index] = std::move(chunk);
+    return chunks_cache_[index];
+}
+
+Result<std::unique_ptr<ChunkedDictionary::Appender>> ChunkedDictionary::Appender::Create(
+    const std::shared_ptr<KeyFactory>& key_factory, int32_t chunk_size_bytes,
+    const std::shared_ptr<MemoryPool>& pool) {
+    return std::unique_ptr<Appender>(new Appender(key_factory, chunk_size_bytes, pool));
+}
+
+ChunkedDictionary::Appender::Appender(const std::shared_ptr<KeyFactory>& key_factory,
+                                      int32_t chunk_size_bytes,
+                                      const std::shared_ptr<MemoryPool>& pool)
+    : pool_(pool),
+      key_factory_(key_factory),
+      chunk_size_bytes_(chunk_size_bytes),
+      chunk_(nullptr),
+      size_(0),
+      key_offset_(0),
+      chunks_offset_(0) {
+    chunks_output_ = std::make_unique<MemorySegmentOutputStream>(
+        MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
+    keys_output_ = std::make_unique<MemorySegmentOutputStream>(
+        MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
+    offsets_output_ = std::make_unique<MemorySegmentOutputStream>(
+        MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
+}
+
+Status ChunkedDictionary::Appender::AppendSorted(const Literal& key, int32_t code) {
+    if (key.IsNull()) {
+        return Status::Invalid("key should not be null");
+    }
+    if (last_key_.has_value()) {
+        PAIMON_ASSIGN_OR_RAISE(int32_t compare_result,
+                               key_factory_->CompareLiteral(*last_key_, key));
+        if (compare_result >= 0) {
+            return Status::Invalid("key must be in sorted order");
+        }
+    }
+    if (last_code_.has_value() && code != *last_code_ + 1) {
+        return Status::Invalid("code must be in sorted order and increment by one");
+    }
+    last_key_ = key;
+    last_code_ = code;
+    if (chunk_ == nullptr) {
+        PAIMON_ASSIGN_OR_RAISE(chunk_,
+                               key_factory_->CreateChunk(key, code, chunk_size_bytes_, pool_));
+    } else {
+        PAIMON_ASSIGN_OR_RAISE(bool success, chunk_->TryAdd(key));
+        if (success) {
+            return Status::OK();
+        }
+        PAIMON_RETURN_NOT_OK(Flush());
+        PAIMON_ASSIGN_OR_RAISE(chunk_,
+                               key_factory_->CreateChunk(key, code, chunk_size_bytes_, pool_));
+    }
+    return Status::OK();
+}
+
+Result<PAIMON_UNIQUE_PTR<Bytes>> ChunkedDictionary::Appender::Serialize() {
+    if (chunk_ != nullptr) {
+        PAIMON_RETURN_NOT_OK(Flush());
+    }
+    int32_t header_size = 0;
+    header_size += sizeof(int8_t);   // version
+    header_size += sizeof(int32_t);  // size
+    header_size += sizeof(int32_t);  // offsets length
+    header_size += sizeof(int32_t);  // chunks length
+    auto data_out = std::make_unique<MemorySegmentOutputStream>(
+        MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_);
+    data_out->WriteValue<int32_t>(header_size);
+    data_out->WriteValue<int8_t>(kCurrentVersion);
+    data_out->WriteValue<int32_t>(size_);
+    data_out->WriteValue<int32_t>(static_cast<int32_t>(offsets_output_->CurrentSize()));
+    data_out->WriteValue<int32_t>(static_cast<int32_t>(chunks_output_->CurrentSize()));
+    PAIMON_RETURN_NOT_OK(MemorySegmentUtils::CopyToStream(
+        offsets_output_->Segments(), 0, static_cast<int32_t>(offsets_output_->CurrentSize()),
+        data_out.get()));
+    PAIMON_RETURN_NOT_OK(MemorySegmentUtils::CopyToStream(
+        chunks_output_->Segments(), 0, static_cast<int32_t>(chunks_output_->CurrentSize()),
+        data_out.get()));
+    PAIMON_RETURN_NOT_OK(MemorySegmentUtils::CopyToStream(
+        keys_output_->Segments(), 0, static_cast<int32_t>(keys_output_->CurrentSize()),
+        data_out.get()));
+    return MemorySegmentUtils::CopyToBytes(
+        data_out->Segments(), 0, static_cast<int32_t>(data_out->CurrentSize()), pool_.get());
+}
+
+Status ChunkedDictionary::Appender::Flush() {
+    chunk_->SetOffset(key_offset_);
+    PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> chunks_bytes, chunk_->SerializeChunk());
+    PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> keys_bytes, chunk_->SerializeKeys());
+    offsets_output_->WriteValue<int32_t>(chunks_offset_);
+    chunks_offset_ += static_cast<int32_t>(chunks_bytes->size());
+    key_offset_ += static_cast<int32_t>(keys_bytes->size());
+    chunks_output_->Write(chunks_bytes->data(), chunks_bytes->size());
+    keys_output_->Write(keys_bytes->data(), keys_bytes->size());
+    size_ += 1;
+    chunk_ = nullptr;
+    return Status::OK();
+}
+
+Result<std::unique_ptr<ChunkedDictionary>> ChunkedDictionary::Create(
+    FieldType field_type, const std::shared_ptr<InputStream>& input_stream, int32_t offset,
+    const std::shared_ptr<MemoryPool>& pool) {
+    auto data_in = std::make_unique<DataInputStream>(input_stream);
+    PAIMON_RETURN_NOT_OK(data_in->Seek(offset));
+    PAIMON_ASSIGN_OR_RAISE(int32_t header_length, data_in->ReadValue<int32_t>());
+    PAIMON_ASSIGN_OR_RAISE(int8_t version, data_in->ReadValue<int8_t>());
+    if (version != kCurrentVersion) {
+        return Status::Invalid("Unknown version of ChunkedDictionary");
+    }
+    PAIMON_ASSIGN_OR_RAISE(int32_t size, data_in->ReadValue<int32_t>());
+    PAIMON_ASSIGN_OR_RAISE(int32_t offsets_length, data_in->ReadValue<int32_t>());
+    PAIMON_ASSIGN_OR_RAISE(int32_t chunks_length, data_in->ReadValue<int32_t>());
+    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeyFactory> factory_shared,
+                           KeyFactory::Create(field_type));
+    auto result = std::unique_ptr<ChunkedDictionary>(new ChunkedDictionary(
+        input_stream, factory_shared, size, offsets_length, chunks_length,
+        static_cast<int32_t>(offset + header_length + sizeof(int32_t)), pool));
+    return result;
+}
+
+ChunkedDictionary::ChunkedDictionary(const std::shared_ptr<InputStream>& input_stream,
+                                     const std::shared_ptr<KeyFactory>& factory, int32_t size,
+                                     int32_t offsets_length, int32_t chunks_length,
+                                     int32_t body_offset, const std::shared_ptr<MemoryPool>& pool)
+    : pool_(pool),
+      factory_(factory),
+      input_stream_(input_stream),
+      size_(size),
+      offsets_length_(offsets_length),
+      chunks_length_(chunks_length),
+      body_offset_(body_offset),
+      offsets_bytes_(nullptr),
+      chunks_bytes_(nullptr),
+      chunks_cache_(std::vector<std::shared_ptr<Chunk>>(size)) {}
+}  // namespace paimon