diff --git a/src/paimon/common/global_index/btree/btree_defs.h b/src/paimon/common/global_index/btree/btree_defs.h new file mode 100644 index 0000000..7521925 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_defs.h @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 +#include + +namespace paimon { +struct BtreeDefs { + BtreeDefs() = delete; + ~BtreeDefs() = delete; + static inline const char kIdentifier[] = "btree"; + /// "btree-index.compression" - The compression algorithm to use for BTreeIndex. + /// Default value is "none". + static inline const char kBtreeIndexCompression[] = "btree-index.compression"; + /// "btree-index.compression-level" - The compression level of the compression algorithm. + /// Default value is 1. + static inline const char kBtreeIndexCompressionLevel[] = "btree-index.compression-level"; + /// "btree-index.block-size" - The block size to use for BTreeIndex. + /// Default value is 64 KB. + static inline const char kBtreeIndexBlockSize[] = "btree-index.block-size"; + /// "btree-index.cache-size" - The cache size to use for BTreeIndex. + /// Default value is 128 MB. + static inline const char kBtreeIndexCacheSize[] = "btree-index.cache-size"; + /// "btree-index.high-priority-pool-ratio" - The high priority pool ratio to use for BTreeIndex. + /// Default value is 0.1. + static inline const char kBtreeIndexHighPriorityPoolRatio[] = + "btree-index.high-priority-pool-ratio"; + + /// "btree-index.read-buffer-size" - Optional. Specifies the read buffer size for the B-tree + /// index. This setting can be tuned based on query patterns: + /// - For range queries (e.g., `VisitLessThan`, `VisitGreaterOrEqual`), increasing the buffer + /// size (e.g., to 1MB) may improve I/O bandwidth and sequential read performance. + /// - For point queries (e.g., `VisitEqual`), buffering can introduce negative effects due to + /// read amplification; it is recommended to leave this option unset. + /// + /// If specified, read block with `BufferedInputStream`. + static inline const char kBtreeIndexReadBufferSize[] = "btree-index.read-buffer-size"; + + static inline const char kDefaultBtreeIndexBlockSize[] = "64KB"; + static inline const char kDefaultBtreeIndexCompression[] = "none"; + static inline const char kDefaultBtreeIndexCacheSize[] = "128MB"; + static inline const int32_t kDefaultBtreeIndexCompressionLevel = 1; + static inline const double kDefaultBtreeIndexHighPriorityPoolRatio = 0.1; +}; +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_file_footer.cpp b/src/paimon/common/global_index/btree/btree_file_footer.cpp new file mode 100644 index 0000000..7b95ace --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_footer.cpp @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_file_footer.h" + +#include + +namespace paimon { + +Result> BTreeFileFooter::Read(MemorySliceInput* input) { + // read version and verify magic number + PAIMON_RETURN_NOT_OK(input->SetPosition(kEncodingLength - 8)); + + int32_t version = input->ReadInt(); + int32_t magic_number = input->ReadInt(); + if (magic_number != kMagicNumber) { + return Status::Invalid( + fmt::format("File is not a btree index file (expected magic number {:#x}, got {:#x})", + kMagicNumber, magic_number)); + } + + PAIMON_RETURN_NOT_OK(input->SetPosition(0)); + + // read bloom filter and index handles + auto offset = input->ReadLong(); + auto size = input->ReadInt(); + auto expected_entries = input->ReadLong(); + std::optional bloom_filter_handle = + BloomFilterHandle(offset, size, expected_entries); + if (bloom_filter_handle->Offset() == 0 && bloom_filter_handle->Size() == 0 && + bloom_filter_handle->ExpectedEntries() == 0) { + bloom_filter_handle = std::nullopt; + } + + offset = input->ReadLong(); + size = input->ReadInt(); + BlockHandle index_block_handle(offset, size); + + offset = input->ReadLong(); + size = input->ReadInt(); + std::optional null_bitmap_handle = BlockHandle(offset, size); + if (null_bitmap_handle->Offset() == 0 && null_bitmap_handle->Size() == 0) { + null_bitmap_handle = std::nullopt; + } + + return std::make_shared(version, bloom_filter_handle, index_block_handle, + null_bitmap_handle); +} + +MemorySlice BTreeFileFooter::Write(const std::shared_ptr& footer, + MemoryPool* pool) { + MemorySliceOutput output(kEncodingLength, pool); + return BTreeFileFooter::Write(footer, &output); +} + +MemorySlice BTreeFileFooter::Write(const std::shared_ptr& footer, + MemorySliceOutput* output) { + // write bloom filter handle + const auto& bloom_filter_handle = footer->GetBloomFilterHandle(); + if (!bloom_filter_handle.has_value()) { + output->WriteValue(static_cast(0)); + output->WriteValue(static_cast(0)); + output->WriteValue(static_cast(0)); + } else { + output->WriteValue(bloom_filter_handle->Offset()); + output->WriteValue(bloom_filter_handle->Size()); + output->WriteValue(bloom_filter_handle->ExpectedEntries()); + } + + // write index block handle + const auto& index_block_handle = footer->GetIndexBlockHandle(); + output->WriteValue(index_block_handle.Offset()); + output->WriteValue(index_block_handle.Size()); + + // write null bitmap handle + const auto& null_bitmap_handle = footer->GetNullBitmapHandle(); + if (!null_bitmap_handle.has_value()) { + output->WriteValue(static_cast(0)); + output->WriteValue(static_cast(0)); + } else { + output->WriteValue(null_bitmap_handle->Offset()); + output->WriteValue(null_bitmap_handle->Size()); + } + + // write version and magic number + output->WriteValue(footer->GetVersion()); + output->WriteValue(kMagicNumber); + + return output->ToSlice(); +} + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_file_footer.h b/src/paimon/common/global_index/btree/btree_file_footer.h new file mode 100644 index 0000000..8364105 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_footer.h @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + +#include "paimon/common/memory/memory_slice_input.h" +#include "paimon/common/memory/memory_slice_output.h" +#include "paimon/common/sst/block_handle.h" +#include "paimon/common/sst/bloom_filter_handle.h" + +namespace paimon { +/// The Footer for BTree file. +class BTreeFileFooter { + public: + static Result> Read(MemorySliceInput* input); + static MemorySlice Write(const std::shared_ptr& footer, MemoryPool* pool); + static MemorySlice Write(const std::shared_ptr& footer, + MemorySliceOutput* output); + + public: + BTreeFileFooter(const std::optional& bloom_filter_handle, + const BlockHandle& index_block_handle, + const std::optional& null_bitmap_handle) + : BTreeFileFooter(kCurrentVersion, bloom_filter_handle, index_block_handle, + null_bitmap_handle) {} + + BTreeFileFooter(int32_t version, const std::optional& bloom_filter_handle, + const BlockHandle& index_block_handle, + const std::optional& null_bitmap_handle) + : version_(version), + bloom_filter_handle_(bloom_filter_handle), + index_block_handle_(index_block_handle), + null_bitmap_handle_(null_bitmap_handle) {} + + int32_t GetVersion() const { + return version_; + } + + const std::optional& GetBloomFilterHandle() const { + return bloom_filter_handle_; + } + + const BlockHandle& GetIndexBlockHandle() const { + return index_block_handle_; + } + + const std::optional& GetNullBitmapHandle() const { + return null_bitmap_handle_; + } + + public: + static constexpr int32_t kMagicNumber = 0x50425449; + static constexpr int32_t kCurrentVersion = 1; + static constexpr int32_t kEncodingLength = 52; + + private: + int32_t version_; + std::optional bloom_filter_handle_; + BlockHandle index_block_handle_; + std::optional null_bitmap_handle_; +}; + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_file_footer_test.cpp b/src/paimon/common/global_index/btree/btree_file_footer_test.cpp new file mode 100644 index 0000000..9665bda --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_footer_test.cpp @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_file_footer.h" + +#include + +#include "paimon/common/sst/block_handle.h" +#include "paimon/common/sst/bloom_filter_handle.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class BTreeFileFooterTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(BTreeFileFooterTest, ReadWriteRoundTrip) { + BloomFilterHandle bloom_filter_handle(100, 50, 1000); + BlockHandle index_block_handle(200, 80); + BlockHandle null_bitmap_handle(300, 40); + + auto footer = std::make_shared(bloom_filter_handle, index_block_handle, + null_bitmap_handle); + + auto serialized = BTreeFileFooter::Write(footer, pool_.get()); + ASSERT_EQ(serialized.Length(), BTreeFileFooter::kEncodingLength); + + auto input = serialized.ToInput(); + ASSERT_OK_AND_ASSIGN(auto deserialized_footer, BTreeFileFooter::Read(&input)); + + const auto& bf_handle = deserialized_footer->GetBloomFilterHandle(); + ASSERT_TRUE(bf_handle.has_value()); + ASSERT_EQ(bf_handle->Offset(), 100); + ASSERT_EQ(bf_handle->Size(), 50); + ASSERT_EQ(bf_handle->ExpectedEntries(), 1000); + + const auto& ib_handle = deserialized_footer->GetIndexBlockHandle(); + ASSERT_EQ(ib_handle.Offset(), 200); + ASSERT_EQ(ib_handle.Size(), 80); + + const auto& nb_handle = deserialized_footer->GetNullBitmapHandle(); + ASSERT_TRUE(nb_handle.has_value()); + ASSERT_EQ(nb_handle->Offset(), 300); + ASSERT_EQ(nb_handle->Size(), 40); +} + +TEST_F(BTreeFileFooterTest, ReadWriteWithNullBloomFilter) { + BlockHandle index_block_handle(200, 80); + BlockHandle null_bitmap_handle(300, 40); + + auto footer = + std::make_shared(std::nullopt, index_block_handle, null_bitmap_handle); + + auto serialized = BTreeFileFooter::Write(footer, pool_.get()); + ASSERT_EQ(serialized.Length(), BTreeFileFooter::kEncodingLength); + + auto input = serialized.ToInput(); + ASSERT_OK_AND_ASSIGN(auto deserialized_footer, BTreeFileFooter::Read(&input)); + + ASSERT_FALSE(deserialized_footer->GetBloomFilterHandle().has_value()); + + const auto& ib_handle = deserialized_footer->GetIndexBlockHandle(); + ASSERT_EQ(ib_handle.Offset(), 200); + ASSERT_EQ(ib_handle.Size(), 80); + + const auto& nb_handle = deserialized_footer->GetNullBitmapHandle(); + ASSERT_TRUE(nb_handle.has_value()); + ASSERT_EQ(nb_handle->Offset(), 300); + ASSERT_EQ(nb_handle->Size(), 40); +} + +TEST_F(BTreeFileFooterTest, ReadWriteWithNullNullBitmap) { + BloomFilterHandle bloom_filter_handle(100, 50, 1000); + BlockHandle index_block_handle(200, 80); + + auto footer = + std::make_shared(bloom_filter_handle, index_block_handle, std::nullopt); + + auto serialized = BTreeFileFooter::Write(footer, pool_.get()); + ASSERT_EQ(serialized.Length(), BTreeFileFooter::kEncodingLength); + + auto input = serialized.ToInput(); + ASSERT_OK_AND_ASSIGN(auto deserialized_footer, BTreeFileFooter::Read(&input)); + + const auto& bf_handle = deserialized_footer->GetBloomFilterHandle(); + ASSERT_TRUE(bf_handle.has_value()); + ASSERT_EQ(bf_handle->Offset(), 100); + ASSERT_EQ(bf_handle->Size(), 50); + ASSERT_EQ(bf_handle->ExpectedEntries(), 1000); + + const auto& ib_handle = deserialized_footer->GetIndexBlockHandle(); + ASSERT_EQ(ib_handle.Offset(), 200); + ASSERT_EQ(ib_handle.Size(), 80); + + ASSERT_FALSE(deserialized_footer->GetNullBitmapHandle().has_value()); +} + +TEST_F(BTreeFileFooterTest, ReadWriteWithAllNullHandles) { + BlockHandle index_block_handle(200, 80); + + auto footer = std::make_shared(std::nullopt, index_block_handle, std::nullopt); + + auto serialized = BTreeFileFooter::Write(footer, pool_.get()); + ASSERT_EQ(serialized.Length(), BTreeFileFooter::kEncodingLength); + + auto input = serialized.ToInput(); + ASSERT_OK_AND_ASSIGN(auto deserialized_footer, BTreeFileFooter::Read(&input)); + + ASSERT_FALSE(deserialized_footer->GetBloomFilterHandle().has_value()); + + const auto& ib_handle = deserialized_footer->GetIndexBlockHandle(); + ASSERT_EQ(ib_handle.Offset(), 200); + ASSERT_EQ(ib_handle.Size(), 80); + + ASSERT_FALSE(deserialized_footer->GetNullBitmapHandle().has_value()); +} + +TEST_F(BTreeFileFooterTest, InvalidMagicNumber) { + MemorySliceOutput output(BTreeFileFooter::kEncodingLength, pool_.get()); + + output.WriteValue(static_cast(0)); + output.WriteValue(static_cast(0)); + output.WriteValue(static_cast(0)); + + output.WriteValue(static_cast(200)); + output.WriteValue(static_cast(80)); + + output.WriteValue(static_cast(0)); + output.WriteValue(static_cast(0)); + + output.WriteValue(static_cast(1)); // version + output.WriteValue(static_cast(12345)); // Invalid magic number + + auto serialized = output.ToSlice(); + auto input = serialized.ToInput(); + + auto deserialized = BTreeFileFooter::Read(&input); + ASSERT_NOK_WITH_MSG(deserialized, "File is not a btree index file (expected magic number"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/global_index/btree/btree_file_meta_selector.cpp b/src/paimon/common/global_index/btree/btree_file_meta_selector.cpp new file mode 100644 index 0000000..0753794 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector.cpp @@ -0,0 +1,192 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_file_meta_selector.h" + +#include "paimon/common/memory/memory_slice.h" + +namespace paimon { +BTreeFileMetaSelector::BTreeFileMetaSelector(const std::vector& files, + const std::shared_ptr& key_type, + const std::shared_ptr& pool) + : key_type_(key_type), + pool_(pool), + comparator_(KeySerializer::CreateComparator(key_type, pool)) { + files_.reserve(files.size()); + for (const auto& file : files) { + auto index_meta = BTreeIndexMeta::Deserialize(file.metadata, pool.get()); + files_.emplace_back(file, std::move(index_meta)); + } +} + +Result> BTreeFileMetaSelector::VisitIsNotNull() { + return Filter([](const BTreeIndexMeta& meta) -> Result { return !meta.OnlyNulls(); }); +} + +Result> BTreeFileMetaSelector::VisitIsNull() { + return Filter([](const BTreeIndexMeta& meta) -> Result { return meta.HasNulls(); }); +} + +Result> BTreeFileMetaSelector::VisitEqual(const Literal& literal) { + PAIMON_ASSIGN_OR_RAISE(MemorySlice literal_slice, SerializeLiteral(literal)); + return Filter([this, &literal_slice](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice min_key_slice = WrapKeySlice(meta.FirstKey()); + MemorySlice max_key_slice = WrapKeySlice(meta.LastKey()); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_min, comparator_(literal_slice, min_key_slice)); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_max, comparator_(literal_slice, max_key_slice)); + return cmp_min >= 0 && cmp_max <= 0; + }); +} + +Result> BTreeFileMetaSelector::VisitNotEqual( + const Literal& literal) { + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::VisitLessThan( + const Literal& literal) { + // file.minKey < literal + PAIMON_ASSIGN_OR_RAISE(MemorySlice literal_slice, SerializeLiteral(literal)); + return Filter([this, &literal_slice](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice min_key_slice = WrapKeySlice(meta.FirstKey()); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, comparator_(min_key_slice, literal_slice)); + return cmp < 0; + }); +} + +Result> BTreeFileMetaSelector::VisitLessOrEqual( + const Literal& literal) { + // file.minKey <= literal + PAIMON_ASSIGN_OR_RAISE(MemorySlice literal_slice, SerializeLiteral(literal)); + return Filter([this, &literal_slice](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice min_key_slice = WrapKeySlice(meta.FirstKey()); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, comparator_(min_key_slice, literal_slice)); + return cmp <= 0; + }); +} + +Result> BTreeFileMetaSelector::VisitGreaterThan( + const Literal& literal) { + // file.maxKey > literal + PAIMON_ASSIGN_OR_RAISE(MemorySlice literal_slice, SerializeLiteral(literal)); + return Filter([this, &literal_slice](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice max_key_slice = WrapKeySlice(meta.LastKey()); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, comparator_(max_key_slice, literal_slice)); + return cmp > 0; + }); +} + +Result> BTreeFileMetaSelector::VisitGreaterOrEqual( + const Literal& literal) { + // file.maxKey >= literal + PAIMON_ASSIGN_OR_RAISE(MemorySlice literal_slice, SerializeLiteral(literal)); + return Filter([this, &literal_slice](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice max_key_slice = WrapKeySlice(meta.LastKey()); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, comparator_(max_key_slice, literal_slice)); + return cmp >= 0; + }); +} + +Result> BTreeFileMetaSelector::VisitIn( + const std::vector& literals) { + std::vector literal_slices; + literal_slices.reserve(literals.size()); + for (const auto& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(MemorySlice slice, SerializeLiteral(literal)); + literal_slices.push_back(std::move(slice)); + } + return Filter([this, &literal_slices](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + MemorySlice min_key_slice = WrapKeySlice(meta.FirstKey()); + MemorySlice max_key_slice = WrapKeySlice(meta.LastKey()); + for (const auto& literal_slice : literal_slices) { + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_min, comparator_(literal_slice, min_key_slice)); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_max, comparator_(literal_slice, max_key_slice)); + if (cmp_min >= 0 && cmp_max <= 0) { + return true; + } + } + return false; + }); +} + +Result> BTreeFileMetaSelector::VisitNotIn( + const std::vector& literals) { + // Cannot filter any file by NOT IN condition + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::VisitStartsWith( + const Literal& prefix) { + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::VisitEndsWith(const Literal& suffix) { + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::VisitContains( + const Literal& literal) { + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::VisitLike(const Literal& literal) { + return Filter([](const BTreeIndexMeta& meta) -> Result { return true; }); +} + +Result> BTreeFileMetaSelector::Filter( + const MetaPredicate& predicate) const { + std::vector result; + for (const auto& [io_meta, index_meta] : files_) { + PAIMON_ASSIGN_OR_RAISE(bool matched, predicate(*index_meta)); + if (matched) { + result.push_back(io_meta); + } + } + return result; +} + +MemorySlice BTreeFileMetaSelector::WrapKeySlice(const std::shared_ptr& key) { + return MemorySlice::Wrap(MemorySegment::WrapView(key->data(), key->size())); +} + +Result BTreeFileMetaSelector::SerializeLiteral(const Literal& literal) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, + KeySerializer::SerializeKey(literal, key_type_, pool_.get())); + return MemorySlice::Wrap(bytes); +} + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_file_meta_selector.h b/src/paimon/common/global_index/btree/btree_file_meta_selector.h new file mode 100644 index 0000000..405efc9 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector.h @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 +#include +#include + +#include "paimon/common/global_index/btree/btree_index_meta.h" +#include "paimon/common/global_index/btree/key_serializer.h" +#include "paimon/common/memory/memory_slice.h" +#include "paimon/global_index/global_index_io_meta.h" +#include "paimon/predicate/function_visitor.h" + +namespace paimon { + +/// Selects candidate BTree index files based on filter predicates. +class BTreeFileMetaSelector : public FunctionVisitor> { + public: + BTreeFileMetaSelector(const std::vector& files, + const std::shared_ptr& key_type, + const std::shared_ptr& pool); + + Result> VisitIsNotNull() override; + Result> VisitIsNull() override; + Result> VisitEqual(const Literal& literal) override; + Result> VisitNotEqual(const Literal& literal) override; + Result> VisitLessThan(const Literal& literal) override; + Result> VisitLessOrEqual(const Literal& literal) override; + Result> VisitGreaterThan(const Literal& literal) override; + Result> VisitGreaterOrEqual(const Literal& literal) override; + Result> VisitIn(const std::vector& literals) override; + Result> VisitNotIn( + const std::vector& literals) override; + Result> VisitStartsWith(const Literal& prefix) override; + Result> VisitEndsWith(const Literal& suffix) override; + Result> VisitContains(const Literal& literal) override; + Result> VisitLike(const Literal& literal) override; + + private: + using MetaPredicate = std::function(const BTreeIndexMeta&)>; + + Result> Filter(const MetaPredicate& predicate) const; + + Result SerializeLiteral(const Literal& literal) const; + + /// Create a non-owning MemorySlice view over the raw bytes of a key, + /// avoiding shared_ptr reference-count overhead. + static MemorySlice WrapKeySlice(const std::shared_ptr& key); + + std::vector>> files_; + std::shared_ptr key_type_; + std::shared_ptr pool_; + MemorySlice::SliceComparator comparator_; +}; + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_file_meta_selector_test.cpp b/src/paimon/common/global_index/btree/btree_file_meta_selector_test.cpp new file mode 100644 index 0000000..4e532fe --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector_test.cpp @@ -0,0 +1,231 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_file_meta_selector.h" + +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/common/global_index/btree/btree_index_meta.h" +#include "paimon/common/global_index/btree/key_serializer.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class BTreeFileMetaSelectorTest : public ::testing::Test { + public: + void SetUp() override { + pool_ = GetDefaultPool(); + key_type_ = arrow::int32(); + + // file1: keys [1, 10], has_nulls=true + // file2: keys [15, 20], has_nulls=false + // file3: keys [21, 30], has_nulls=true + // file4: keys [1, 5], has_nulls=false + // file5: keys [19, 25], has_nulls=true + auto meta1 = std::make_shared(SerializeInt(1), SerializeInt(10), true); + auto meta2 = std::make_shared(SerializeInt(15), SerializeInt(20), false); + auto meta3 = std::make_shared(SerializeInt(21), SerializeInt(30), true); + auto meta4 = std::make_shared(SerializeInt(1), SerializeInt(5), false); + auto meta5 = std::make_shared(SerializeInt(19), SerializeInt(25), true); + + // file6: only-nulls file (no keys, has_nulls=true) + auto meta6 = std::make_shared(nullptr, nullptr, true); + files_ = { + GlobalIndexIOMeta("file1", 1, meta1->Serialize(pool_.get())), + GlobalIndexIOMeta("file2", 1, meta2->Serialize(pool_.get())), + GlobalIndexIOMeta("file3", 1, meta3->Serialize(pool_.get())), + GlobalIndexIOMeta("file4", 1, meta4->Serialize(pool_.get())), + GlobalIndexIOMeta("file5", 1, meta5->Serialize(pool_.get())), + GlobalIndexIOMeta("file6", 1, meta6->Serialize(pool_.get())), + }; + } + + std::shared_ptr SerializeInt(int32_t value) const { + Literal literal(value); + EXPECT_OK_AND_ASSIGN(std::shared_ptr result, + KeySerializer::SerializeKey(literal, key_type_, pool_.get())); + return result; + } + + std::set FileNames(const std::vector& metas) const { + std::set names; + for (const auto& meta : metas) { + names.insert(meta.file_path); + } + return names; + } + + void CheckResult(const std::vector& actual, + const std::set& expected) const { + ASSERT_EQ(FileNames(actual), expected); + } + + private: + std::shared_ptr pool_; + std::shared_ptr key_type_; + std::vector files_; +}; + +TEST_F(BTreeFileMetaSelectorTest, TestVisitLessThan) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // minKey < 8: file1(1), file4(1) + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitLessThan(Literal(8))); + CheckResult(result, {"file1", "file4"}); + + // minKey < 15: file1(1), file4(1) (file2 minKey=15, not < 15) + ASSERT_OK_AND_ASSIGN(result, selector.VisitLessThan(Literal(15))); + CheckResult(result, {"file1", "file4"}); + + // minKey < 1: no file has minKey < 1 + ASSERT_OK_AND_ASSIGN(result, selector.VisitLessThan(Literal(1))); + ASSERT_TRUE(result.empty()); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitLessOrEqual) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // minKey <= 20: file1(1), file2(15), file4(1), file5(19) + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitLessOrEqual(Literal(20))); + CheckResult(result, {"file1", "file2", "file4", "file5"}); + + // minKey <= 15: file1(1), file2(15), file4(1) + ASSERT_OK_AND_ASSIGN(result, selector.VisitLessOrEqual(Literal(15))); + CheckResult(result, {"file1", "file2", "file4"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitGreaterThan) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // maxKey > 20: file3(30), file5(25) + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitGreaterThan(Literal(20))); + CheckResult(result, {"file3", "file5"}); + + // maxKey > 30: no file + ASSERT_OK_AND_ASSIGN(result, selector.VisitGreaterThan(Literal(30))); + ASSERT_TRUE(result.empty()); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitGreaterOrEqual) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // maxKey >= 5: all non-null files (file1..file5) + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitGreaterOrEqual(Literal(5))); + CheckResult(result, {"file1", "file2", "file3", "file4", "file5"}); + + // maxKey >= 20: file2(20), file3(30), file5(25) + ASSERT_OK_AND_ASSIGN(result, selector.VisitGreaterOrEqual(Literal(20))); + CheckResult(result, {"file2", "file3", "file5"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitEqual) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // 22 in [21,30] and [19,25] + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitEqual(Literal(22))); + CheckResult(result, {"file3", "file5"}); + + // 30 in [21,30] only + ASSERT_OK_AND_ASSIGN(result, selector.VisitEqual(Literal(30))); + CheckResult(result, {"file3"}); + + // 100 out of all ranges + ASSERT_OK_AND_ASSIGN(result, selector.VisitEqual(Literal(100))); + ASSERT_TRUE(result.empty()); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitNotEqual) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // NotEqual cannot prune any file, returns all + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitNotEqual(Literal(22))); + CheckResult(result, {"file1", "file2", "file3", "file4", "file5", "file6"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitIsNull) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // has_nulls: file1, file3, file5, file6 + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitIsNull()); + CheckResult(result, {"file1", "file3", "file5", "file6"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitIsNotNull) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // !onlyNulls: file1..file5 (file6 is only-nulls) + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitIsNotNull()); + CheckResult(result, {"file1", "file2", "file3", "file4", "file5"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitIn) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // IN(1, 2, 3, 26, 27, 28): + // 1 in [1,10]=file1, [1,5]=file4 + // 2 in [1,10]=file1, [1,5]=file4 + // 3 in [1,10]=file1, [1,5]=file4 + // 26 in [21,30]=file3 + // 27 in [21,30]=file3 + // 28 in [21,30]=file3 + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitIn({Literal(1), Literal(2), Literal(3), + Literal(26), Literal(27), Literal(28)})); + CheckResult(result, {"file1", "file3", "file4"}); + + // IN(100): no match + ASSERT_OK_AND_ASSIGN(result, selector.VisitIn({Literal(100)})); + ASSERT_TRUE(result.empty()); +} + +TEST_F(BTreeFileMetaSelectorTest, TestVisitNotIn) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // NotIn cannot prune any file + ASSERT_OK_AND_ASSIGN(auto result, + selector.VisitNotIn({Literal(1), Literal(7), Literal(19), Literal(30)})); + CheckResult(result, {"file1", "file2", "file3", "file4", "file5", "file6"}); +} + +TEST_F(BTreeFileMetaSelectorTest, TestOnlyNullsFileExcludedFromRangeQueries) { + BTreeFileMetaSelector selector(files_, key_type_, pool_); + + // file6 is only-nulls, should be excluded from all range/equality queries + ASSERT_OK_AND_ASSIGN(auto result, selector.VisitEqual(Literal(1))); + auto names = FileNames(result); + ASSERT_EQ(names.count("file6"), 0u); + + ASSERT_OK_AND_ASSIGN(result, selector.VisitLessThan(Literal(100))); + names = FileNames(result); + ASSERT_EQ(names.count("file6"), 0u); + + ASSERT_OK_AND_ASSIGN(result, selector.VisitGreaterThan(Literal(0))); + names = FileNames(result); + ASSERT_EQ(names.count("file6"), 0u); + + // But IsNull should include file6 + ASSERT_OK_AND_ASSIGN(result, selector.VisitIsNull()); + names = FileNames(result); + ASSERT_EQ(names.count("file6"), 1u); +} + +} // namespace paimon::test diff --git a/src/paimon/common/global_index/btree/btree_index_meta.cpp b/src/paimon/common/global_index/btree/btree_index_meta.cpp new file mode 100644 index 0000000..94d7b44 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_index_meta.cpp @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_index_meta.h" + +#include "paimon/common/memory/memory_slice_output.h" + +namespace paimon { + +std::shared_ptr BTreeIndexMeta::Deserialize(const std::shared_ptr& meta, + paimon::MemoryPool* pool) { + auto slice = MemorySlice::Wrap(meta); + auto input = slice.ToInput(); + auto first_key_len = input.ReadInt(); + std::shared_ptr first_key; + if (first_key_len) { + first_key = input.ReadSliceView(first_key_len).CopyBytes(pool); + } + auto last_key_len = input.ReadInt(); + std::shared_ptr last_key; + if (last_key_len) { + last_key = input.ReadSliceView(last_key_len).CopyBytes(pool); + } + auto has_nulls = input.ReadByte() == static_cast(1); + return std::make_shared(first_key, last_key, has_nulls); +} + +std::shared_ptr BTreeIndexMeta::Serialize(paimon::MemoryPool* pool) const { + // Calculate total size: first_key_len(4) + first_key + last_key_len(4) + last_key + + // has_nulls(1) + int32_t first_key_size = first_key_ ? first_key_->size() : 0; + int32_t last_key_size = last_key_ ? last_key_->size() : 0; + int32_t total_size = Size(); + + MemorySliceOutput output(total_size, pool); + + // Write first_key_len and first_key + output.WriteValue(first_key_size); + if (first_key_) { + output.WriteBytes(first_key_); + } + + // Write last_key_len and last_key + output.WriteValue(last_key_size); + if (last_key_) { + output.WriteBytes(last_key_); + } + + // Write has_nulls + output.WriteValue(static_cast(has_nulls_ ? 1 : 0)); + + return output.ToSlice().GetOrCreateHeapMemory(pool); +} + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_index_meta.h b/src/paimon/common/global_index/btree/btree_index_meta.h new file mode 100644 index 0000000..05f08f9 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_index_meta.h @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + +#include "paimon/common/memory/memory_slice_input.h" +#include "paimon/memory/bytes.h" + +namespace paimon { +/// Index Meta of each BTree index file. The first key and last key of this meta could be null if +/// the entire btree index file only contains nulls. +class BTreeIndexMeta { + public: + static std::shared_ptr Deserialize(const std::shared_ptr& meta, + paimon::MemoryPool* pool); + std::shared_ptr Serialize(paimon::MemoryPool* pool) const; + + public: + BTreeIndexMeta(const std::shared_ptr& first_key, const std::shared_ptr& last_key, + bool has_nulls) + : first_key_(first_key), last_key_(last_key), has_nulls_(has_nulls) {} + + const std::shared_ptr& FirstKey() const { + return first_key_; + } + + const std::shared_ptr& LastKey() const { + return last_key_; + } + + bool HasNulls() const { + return has_nulls_; + } + + bool OnlyNulls() const { + return !(first_key_ || last_key_); + } + + private: + int32_t Size() const { + // 9 bytes => first_key_len(4 byte) + last_key_len(4 byte) + has_null(1 byte) + return (first_key_ ? first_key_->size() : 0) + (last_key_ ? last_key_->size() : 0) + 9; + } + + private: + std::shared_ptr first_key_; + std::shared_ptr last_key_; + bool has_nulls_; +}; + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_index_meta_test.cpp b/src/paimon/common/global_index/btree/btree_index_meta_test.cpp new file mode 100644 index 0000000..c29ec6b --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_index_meta_test.cpp @@ -0,0 +1,202 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/btree_index_meta.h" + +#include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon::test { +class BTreeIndexMetaTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeNormalKeys) { + auto first_key = std::make_shared("first_key_data", pool_.get()); + auto last_key = std::make_shared("last_key_data", pool_.get()); + auto meta = std::make_shared(first_key, last_key, true); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + ASSERT_GT(serialized->size(), 0u); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify first_key + auto deserialized_first = deserialized->FirstKey(); + ASSERT_TRUE(deserialized_first); + ASSERT_EQ(std::string(deserialized_first->data(), deserialized_first->size()), + "first_key_data"); + + // Verify last_key + auto deserialized_last = deserialized->LastKey(); + ASSERT_TRUE(deserialized_last); + ASSERT_EQ(std::string(deserialized_last->data(), deserialized_last->size()), "last_key_data"); + + // Verify has_nulls + ASSERT_TRUE(deserialized->HasNulls()); +} + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeEmptyKeys) { + // Create a BTreeIndexMeta with empty keys (OnlyNulls case) + auto meta = std::make_shared(nullptr, nullptr, true); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify keys are null + ASSERT_FALSE(deserialized->FirstKey()); + ASSERT_FALSE(deserialized->LastKey()); + + // Verify has_nulls + ASSERT_TRUE(deserialized->HasNulls()); + + // Verify OnlyNulls + ASSERT_TRUE(deserialized->OnlyNulls()); +} + +TEST_F(BTreeIndexMetaTest, HasNullsAndOnlyNulls) { + // Case 1: Has nulls with keys + auto meta1 = + std::make_shared(std::make_shared("key", pool_.get()), + std::make_shared("key", pool_.get()), true); + ASSERT_TRUE(meta1->HasNulls()); + ASSERT_FALSE(meta1->OnlyNulls()); + + // Case 2: No nulls with keys + auto meta2 = + std::make_shared(std::make_shared("key", pool_.get()), + std::make_shared("key", pool_.get()), false); + ASSERT_FALSE(meta2->HasNulls()); + ASSERT_FALSE(meta2->OnlyNulls()); + + // Case 3: Only nulls (no keys) + auto meta3 = std::make_shared(nullptr, nullptr, true); + ASSERT_TRUE(meta3->HasNulls()); + ASSERT_TRUE(meta3->OnlyNulls()); + + // Case 4: No nulls and no keys (edge case) + auto meta4 = std::make_shared(nullptr, nullptr, false); + ASSERT_FALSE(meta4->HasNulls()); + ASSERT_TRUE(meta4->OnlyNulls()); +} + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeNoNulls) { + // Create a BTreeIndexMeta without nulls + auto first_key = std::make_shared("abc", pool_.get()); + auto last_key = std::make_shared("xyz", pool_.get()); + auto meta = std::make_shared(first_key, last_key, false); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify has_nulls is false + ASSERT_FALSE(deserialized->HasNulls()); +} + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeWithOnlyFirstKey) { + // Create a BTreeIndexMeta with only first_key (edge case) + auto first_key = std::make_shared("first", pool_.get()); + auto meta = std::make_shared(first_key, nullptr, false); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify first_key + auto deserialized_first = deserialized->FirstKey(); + ASSERT_TRUE(deserialized_first); + ASSERT_EQ(std::string(deserialized_first->data(), deserialized_first->size()), "first"); + + // Verify last_key is null + ASSERT_FALSE(deserialized->LastKey()); +} + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeWithOnlyLastKey) { + // Create a BTreeIndexMeta with only last_key (edge case) + auto last_key = std::make_shared("last", pool_.get()); + auto meta = std::make_shared(nullptr, last_key, false); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify first_key is null + ASSERT_FALSE(deserialized->FirstKey()); + + // Verify last_key + auto deserialized_last = deserialized->LastKey(); + ASSERT_TRUE(deserialized_last); + ASSERT_EQ(std::string(deserialized_last->data(), deserialized_last->size()), "last"); +} + +TEST_F(BTreeIndexMetaTest, SerializeDeserializeBinaryKeys) { + // Create a BTreeIndexMeta with binary keys containing null bytes + std::string binary_first = std::string("key\0with\0nulls", 14); + std::string binary_last = std::string("last\0key", 8); + auto first_key = std::make_shared(binary_first, pool_.get()); + auto last_key = std::make_shared(binary_last, pool_.get()); + auto meta = std::make_shared(first_key, last_key, true); + + // Serialize + auto serialized = meta->Serialize(pool_.get()); + ASSERT_TRUE(serialized); + + // Deserialize + auto deserialized = BTreeIndexMeta::Deserialize(serialized, pool_.get()); + ASSERT_TRUE(deserialized); + + // Verify first_key + auto deserialized_first = deserialized->FirstKey(); + ASSERT_TRUE(deserialized_first); + ASSERT_EQ(std::string(deserialized_first->data(), deserialized_first->size()), binary_first); + + // Verify last_key + auto deserialized_last = deserialized->LastKey(); + ASSERT_TRUE(deserialized_last); + ASSERT_EQ(std::string(deserialized_last->data(), deserialized_last->size()), binary_last); +} + +} // namespace paimon::test diff --git a/src/paimon/common/global_index/btree/key_serializer.cpp b/src/paimon/common/global_index/btree/key_serializer.cpp new file mode 100644 index 0000000..d6d7175 --- /dev/null +++ b/src/paimon/common/global_index/btree/key_serializer.cpp @@ -0,0 +1,269 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/key_serializer.h" + +#include "fmt/format.h" +#include "paimon/common/memory/memory_slice_input.h" +#include "paimon/common/memory/memory_slice_output.h" +#include "paimon/common/utils/date_time_utils.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/preconditions.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/status.h" +namespace paimon { +Result> KeySerializer::SerializeKey( + const Literal& literal, const std::shared_ptr& type, MemoryPool* pool) { + if (literal.IsNull()) { + return Status::Invalid("cannot serialize null in KeySerializer"); + } + switch (literal.GetType()) { + case FieldType::BOOLEAN: { + MemorySliceOutput output(1, pool); + output.Reset(); + output.WriteValue(literal.GetValue() ? static_cast(1) + : static_cast(0)); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::TINYINT: { + MemorySliceOutput output(1, pool); + output.Reset(); + output.WriteValue(literal.GetValue()); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::SMALLINT: { + MemorySliceOutput output(2, pool); + output.Reset(); + output.WriteValue(literal.GetValue()); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::INT: + case FieldType::DATE: { + MemorySliceOutput output(4, pool); + output.Reset(); + output.WriteValue(literal.GetValue()); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::BIGINT: { + MemorySliceOutput output(8, pool); + output.Reset(); + output.WriteValue(literal.GetValue()); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::FLOAT: { + MemorySliceOutput output(4, pool); + output.Reset(); + auto fvalue = literal.GetValue(); + int32_t ivalue; + memcpy(&ivalue, &fvalue, sizeof(float)); + output.WriteValue(ivalue); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::DOUBLE: { + MemorySliceOutput output(8, pool); + output.Reset(); + auto dvalue = literal.GetValue(); + int64_t ivalue; + memcpy(&ivalue, &dvalue, sizeof(double)); + output.WriteValue(ivalue); + return output.ToSlice().CopyBytes(pool); + } + case FieldType::STRING: { + auto svalue = literal.GetValue(); + std::shared_ptr bytes = Bytes::AllocateBytes(svalue, pool); + return bytes; + } + case FieldType::TIMESTAMP: { + auto ts_type = std::dynamic_pointer_cast(type); + PAIMON_RETURN_NOT_OK(Preconditions::CheckNotNull( + ts_type, "ts type cannot cast to arrow::TimestampType in BTreeGlobalIndex")); + MemorySliceOutput output(8, pool); + output.Reset(); + auto ts = literal.GetValue(); + if (Timestamp::IsCompact(DateTimeUtils::GetPrecisionFromType(ts_type))) { + output.WriteValue(ts.GetMillisecond()); + } else { + output.WriteValue(ts.GetMillisecond()); + PAIMON_RETURN_NOT_OK(output.WriteVarLenInt(ts.GetNanoOfMillisecond())); + } + return output.ToSlice().CopyBytes(pool); + } + case FieldType::DECIMAL: { + auto decimal_type = std::dynamic_pointer_cast(type); + PAIMON_RETURN_NOT_OK(Preconditions::CheckNotNull( + decimal_type, + "decimal type cannot cast to arrow::Decimal128Type in BTreeGlobalIndex")); + + auto decimal = literal.GetValue(); + if (Decimal::IsCompact(decimal_type->precision())) { + MemorySliceOutput output(8, pool); + output.Reset(); + output.WriteValue(decimal.ToUnscaledLong()); + return output.ToSlice().CopyBytes(pool); + } else { + std::vector decimal_bytes = decimal.ToUnscaledBytes(); + std::shared_ptr bytes = Bytes::AllocateBytes(decimal_bytes.size(), pool); + memcpy(bytes->data(), decimal_bytes.data(), decimal_bytes.size()); + return bytes; + } + } + default: + return Status::Invalid( + fmt::format("Not support serialize {} type in BTreeGlobalIndex", + FieldTypeUtils::FieldTypeToString(literal.GetType()))); + } +} + +Result KeySerializer::DeserializeKey(const MemorySlice& slice, + const std::shared_ptr& type, + MemoryPool* pool) { + switch (type->id()) { + case arrow::Type::type::BOOL: + return Literal(slice.ReadByte(0) == 1 ? true : false); + case arrow::Type::type::INT8: + return Literal(slice.ReadByte(0)); + case arrow::Type::type::INT16: + return Literal(slice.ReadShort(0)); + case arrow::Type::type::INT32: + return Literal(slice.ReadInt(0)); + case arrow::Type::type::DATE32: + return Literal(FieldType::DATE, slice.ReadInt(0)); + case arrow::Type::type::INT64: + return Literal(slice.ReadLong(0)); + case arrow::Type::type::FLOAT: { + int32_t ivalue = slice.ReadInt(0); + float fvalue; + memcpy(&fvalue, &ivalue, sizeof(fvalue)); + return Literal(fvalue); + } + case arrow::Type::type::DOUBLE: { + int64_t ivalue = slice.ReadLong(0); + double dvalue; + memcpy(&dvalue, &ivalue, sizeof(dvalue)); + return Literal(dvalue); + } + case arrow::Type::type::STRING: { + auto bytes = slice.CopyBytes(pool); + return Literal(FieldType::STRING, bytes->data(), bytes->size()); + } + case arrow::Type::type::TIMESTAMP: { + auto ts_type = std::dynamic_pointer_cast(type); + PAIMON_RETURN_NOT_OK(Preconditions::CheckNotNull( + ts_type, "ts type cannot cast to arrow::TimestampType in BTreeGlobalIndex")); + if (Timestamp::IsCompact(DateTimeUtils::GetPrecisionFromType(ts_type))) { + return Literal(Timestamp::FromEpochMillis(slice.ReadLong(0))); + } else { + auto input = slice.ToInput(); + int64_t millis = input.ReadLong(); + PAIMON_ASSIGN_OR_RAISE(int32_t nanos, input.ReadVarLenInt()); + return Literal(Timestamp(millis, nanos)); + } + } + case arrow::Type::type::DECIMAL128: { + auto decimal_type = std::dynamic_pointer_cast(type); + PAIMON_RETURN_NOT_OK(Preconditions::CheckNotNull( + decimal_type, + "decimal type cannot cast to arrow::Decimal128Type in BTreeGlobalIndex")); + if (Decimal::IsCompact(decimal_type->precision())) { + return Literal(Decimal::FromUnscaledLong( + slice.ReadLong(0), decimal_type->precision(), decimal_type->scale())); + } else { + auto bytes = slice.CopyBytes(pool); + return Literal(Decimal::FromUnscaledBytes(decimal_type->precision(), + decimal_type->scale(), bytes.get())); + } + } + default: + return Status::Invalid(fmt::format( + "Not support deserialize {} type in BTreeGlobalIndex", type->ToString())); + } +} + +MemorySlice::SliceComparator KeySerializer::CreateComparator( + const std::shared_ptr& type, const std::shared_ptr& pool) { + // Fast paths for integer and string types: direct value comparison without Literal + // deserialization, avoiding heap allocations entirely. + switch (type->id()) { + case arrow::Type::type::STRING: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + std::string_view sv_a = a.ReadStringView(); + std::string_view sv_b = b.ReadStringView(); + int32_t cmp = sv_a.compare(sv_b); + return cmp == 0 ? 0 : (cmp > 0 ? 1 : -1); + }; + case arrow::Type::type::BOOL: + case arrow::Type::type::INT8: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int8_t va = a.ReadByte(0); + int8_t vb = b.ReadByte(0); + return (va < vb) ? -1 : (va > vb ? 1 : 0); + }; + case arrow::Type::type::INT16: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int16_t va = a.ReadShort(0); + int16_t vb = b.ReadShort(0); + return (va < vb) ? -1 : (va > vb ? 1 : 0); + }; + case arrow::Type::type::INT32: + case arrow::Type::type::DATE32: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int32_t va = a.ReadInt(0); + int32_t vb = b.ReadInt(0); + return (va < vb) ? -1 : (va > vb ? 1 : 0); + }; + case arrow::Type::type::INT64: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int64_t va = a.ReadLong(0); + int64_t vb = b.ReadLong(0); + return (va < vb) ? -1 : (va > vb ? 1 : 0); + }; + case arrow::Type::type::FLOAT: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int32_t ia = a.ReadInt(0); + int32_t ib = b.ReadInt(0); + float fa, fb; + memcpy(&fa, &ia, sizeof(float)); + memcpy(&fb, &ib, sizeof(float)); + return FieldsComparator::CompareFloatingPoint(fa, fb); + }; + case arrow::Type::type::DOUBLE: + return [](const MemorySlice& a, const MemorySlice& b) -> Result { + int64_t ia = a.ReadLong(0); + int64_t ib = b.ReadLong(0); + double da, db; + memcpy(&da, &ia, sizeof(double)); + memcpy(&db, &ib, sizeof(double)); + return FieldsComparator::CompareFloatingPoint(da, db); + }; + default: + break; + } + + // Fallback for complex types (TIMESTAMP, DECIMAL, etc.): + // deserialize to Literal and compare. + return + [pool = pool, type = type](const MemorySlice& a, const MemorySlice& b) -> Result { + PAIMON_ASSIGN_OR_RAISE(Literal la, DeserializeKey(a, type, pool.get())); + PAIMON_ASSIGN_OR_RAISE(Literal lb, DeserializeKey(b, type, pool.get())); + return la.CompareTo(lb); + }; +} +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/key_serializer.h b/src/paimon/common/global_index/btree/key_serializer.h new file mode 100644 index 0000000..82f8510 --- /dev/null +++ b/src/paimon/common/global_index/btree/key_serializer.h @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 "arrow/api.h" +#include "paimon/common/memory/memory_slice.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +namespace paimon { +class KeySerializer { + public: + KeySerializer() = delete; + ~KeySerializer() = delete; + + static Result> SerializeKey(const Literal& literal, + const std::shared_ptr& type, + MemoryPool* pool); + + static Result DeserializeKey(const MemorySlice& slice, + const std::shared_ptr& type, + MemoryPool* pool); + + static MemorySlice::SliceComparator CreateComparator( + const std::shared_ptr& type, const std::shared_ptr& pool); +}; +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/key_serializer_test.cpp b/src/paimon/common/global_index/btree/key_serializer_test.cpp new file mode 100644 index 0000000..3dc8a0f --- /dev/null +++ b/src/paimon/common/global_index/btree/key_serializer_test.cpp @@ -0,0 +1,261 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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/global_index/btree/key_serializer.h" + +#include "gtest/gtest.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class KeySerializerTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + } + + std::shared_ptr pool_; +}; + +TEST_F(KeySerializerTest, SerializeAndDeserializeAllTypes) { + // BOOLEAN + { + Literal literal(true); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::boolean(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::boolean(), pool_.get())); + ASSERT_EQ(result.GetValue(), true); + + Literal literal_false(false); + ASSERT_OK_AND_ASSIGN( + bytes, KeySerializer::SerializeKey(literal_false, arrow::boolean(), pool_.get())); + slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(result, + KeySerializer::DeserializeKey(slice, arrow::boolean(), pool_.get())); + ASSERT_EQ(result.GetValue(), false); + } + + // TINYINT (int8) + { + Literal literal(static_cast(-42)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::int8(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::int8(), pool_.get())); + ASSERT_EQ(result.GetValue(), -42); + } + + // SMALLINT (int16) + { + Literal literal(static_cast(12345)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::int16(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::int16(), pool_.get())); + ASSERT_EQ(result.GetValue(), 12345); + } + + // INT (int32) + { + Literal literal(42); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::int32(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::int32(), pool_.get())); + ASSERT_EQ(result.GetValue(), 42); + } + + // DATE (stored as int32) + { + Literal literal(FieldType::DATE, 18000); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::date32(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::date32(), pool_.get())); + ASSERT_EQ(result.GetType(), FieldType::DATE); + ASSERT_EQ(result.GetValue(), 18000); + } + + // BIGINT (int64) + { + Literal literal(static_cast(123456789012345LL)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::int64(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::int64(), pool_.get())); + ASSERT_EQ(result.GetValue(), 123456789012345LL); + } + + // FLOAT + { + Literal literal(3.14f); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::float32(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::float32(), pool_.get())); + ASSERT_FLOAT_EQ(result.GetValue(), 3.14f); + } + + // DOUBLE + { + Literal literal(2.718281828); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::float64(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::float64(), pool_.get())); + ASSERT_DOUBLE_EQ(result.GetValue(), 2.718281828); + } + + // STRING + { + Literal literal(FieldType::STRING, "hello world", 11); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, arrow::utf8(), pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, arrow::utf8(), pool_.get())); + ASSERT_EQ(result.GetValue(), "hello world"); + } + + // TIMESTAMP (compact, millis precision) + { + auto ts_type = arrow::timestamp(arrow::TimeUnit::MILLI); + Literal literal(Timestamp::FromEpochMillis(1234567890)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, ts_type, pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, ts_type, pool_.get())); + ASSERT_EQ(result.GetValue().GetMillisecond(), 1234567890); + } + + // TIMESTAMP (non-compact, nano precision) + { + auto ts_type = arrow::timestamp(arrow::TimeUnit::NANO); + Literal literal(Timestamp(5000, 123456)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, ts_type, pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, ts_type, pool_.get())); + ASSERT_EQ(result.GetValue().GetMillisecond(), 5000); + ASSERT_EQ(result.GetValue().GetNanoOfMillisecond(), 123456); + } + + // DECIMAL (compact, precision <= 18) + { + auto decimal_type = arrow::decimal128(10, 2); + Literal literal(Decimal::FromUnscaledLong(12345, 10, 2)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, decimal_type, pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, decimal_type, pool_.get())); + ASSERT_EQ(result.GetValue().ToUnscaledLong(), 12345); + } + + // DECIMAL (non-compact, precision > 18) + { + auto decimal_type = arrow::decimal128(25, 3); + Literal literal(Decimal(25, 3, 9999999)); + ASSERT_OK_AND_ASSIGN(auto bytes, + KeySerializer::SerializeKey(literal, decimal_type, pool_.get())); + auto slice = MemorySlice::Wrap(MemorySegment::Wrap(bytes)); + ASSERT_OK_AND_ASSIGN(auto result, + KeySerializer::DeserializeKey(slice, decimal_type, pool_.get())); + ASSERT_EQ(result.GetValue().ToString(), literal.GetValue().ToString()); + } + + // NULL should fail + { + Literal null_literal(FieldType::INT); + ASSERT_NOK_WITH_MSG(KeySerializer::SerializeKey(null_literal, arrow::int32(), pool_.get()), + "cannot serialize null in KeySerializer"); + } + + // unsupported type + { + Literal literal(FieldType::BINARY, "ab", 2); + ASSERT_NOK_WITH_MSG(KeySerializer::SerializeKey(literal, arrow::binary(), pool_.get()), + "Not support serialize BINARY type in BTreeGlobalIndex"); + } +} + +TEST_F(KeySerializerTest, CreateComparator) { + // INT comparator + { + auto comparator = KeySerializer::CreateComparator(arrow::int32(), pool_); + Literal literal_1(1); + Literal literal_2(2); + Literal literal_3(1); + ASSERT_OK_AND_ASSIGN(auto bytes_1, + KeySerializer::SerializeKey(literal_1, arrow::int32(), pool_.get())); + ASSERT_OK_AND_ASSIGN(auto bytes_2, + KeySerializer::SerializeKey(literal_2, arrow::int32(), pool_.get())); + ASSERT_OK_AND_ASSIGN(auto bytes_3, + KeySerializer::SerializeKey(literal_3, arrow::int32(), pool_.get())); + auto slice_1 = MemorySlice::Wrap(MemorySegment::Wrap(bytes_1)); + auto slice_2 = MemorySlice::Wrap(MemorySegment::Wrap(bytes_2)); + auto slice_3 = MemorySlice::Wrap(MemorySegment::Wrap(bytes_3)); + + ASSERT_OK_AND_ASSIGN(auto cmp_result, comparator(slice_1, slice_2)); + ASSERT_LT(cmp_result, 0); + ASSERT_OK_AND_ASSIGN(cmp_result, comparator(slice_2, slice_1)); + ASSERT_GT(cmp_result, 0); + ASSERT_OK_AND_ASSIGN(cmp_result, comparator(slice_1, slice_3)); + ASSERT_EQ(cmp_result, 0); + } + + // STRING comparator + { + auto comparator = KeySerializer::CreateComparator(arrow::utf8(), pool_); + Literal literal_a(FieldType::STRING, "apple", 5); + Literal literal_b(FieldType::STRING, "banana", 6); + Literal literal_c(FieldType::STRING, "apple", 5); + ASSERT_OK_AND_ASSIGN(auto bytes_a, + KeySerializer::SerializeKey(literal_a, arrow::utf8(), pool_.get())); + ASSERT_OK_AND_ASSIGN(auto bytes_b, + KeySerializer::SerializeKey(literal_b, arrow::utf8(), pool_.get())); + ASSERT_OK_AND_ASSIGN(auto bytes_c, + KeySerializer::SerializeKey(literal_c, arrow::utf8(), pool_.get())); + auto slice_a = MemorySlice::Wrap(MemorySegment::Wrap(bytes_a)); + auto slice_b = MemorySlice::Wrap(MemorySegment::Wrap(bytes_b)); + auto slice_c = MemorySlice::Wrap(MemorySegment::Wrap(bytes_c)); + + ASSERT_OK_AND_ASSIGN(auto cmp_result, comparator(slice_a, slice_b)); + ASSERT_LT(cmp_result, 0); + ASSERT_OK_AND_ASSIGN(cmp_result, comparator(slice_b, slice_a)); + ASSERT_GT(cmp_result, 0); + ASSERT_OK_AND_ASSIGN(cmp_result, comparator(slice_a, slice_c)); + ASSERT_EQ(cmp_result, 0); + } +} + +} // namespace paimon::test