From 152334aaabce74bcab4632f8c987c89d474a47be Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Fri, 24 Apr 2026 09:42:59 +0000 Subject: [PATCH 1/3] feat: support BTreeFileMetaSelector & LazyFilteredBTreeReader for btree index --- .../global_index/global_index_io_meta.h | 9 +- .../paimon/global_index/global_index_writer.h | 3 +- src/paimon/CMakeLists.txt | 1 + src/paimon/common/global_index/CMakeLists.txt | 2 + .../bitmap/bitmap_global_index.cpp | 4 +- .../bitmap/bitmap_global_index_test.cpp | 17 +- .../btree/btree_compatibility_test.cpp | 222 ++++++++++++-- .../btree/btree_file_meta_selector.cpp | 172 +++++++++++ .../btree/btree_file_meta_selector.h | 65 +++++ .../btree/btree_file_meta_selector_test.cpp | 228 +++++++++++++++ .../btree/btree_global_index_factory.cpp | 2 +- .../btree_global_index_integration_test.cpp | 94 +++--- .../btree/btree_global_index_writer.cpp | 20 +- .../btree/btree_global_index_writer.h | 9 +- .../btree/btree_global_indexer.cpp | 122 ++------ .../global_index/btree/btree_global_indexer.h | 11 +- .../global_index/btree/key_serializer.cpp | 2 - .../btree/lazy_filtered_btree_reader.cpp | 270 ++++++++++++++++++ .../btree/lazy_filtered_btree_reader.h | 98 +++++++ .../common/global_index/global_index_utils.h | 53 ++++ .../rangebitmap/range_bitmap_global_index.cpp | 4 +- .../range_bitmap_global_index_test.cpp | 5 +- .../wrap/file_index_reader_wrapper.h | 9 +- .../wrap/file_index_reader_wrapper_test.cpp | 14 +- .../wrap/file_index_writer_wrapper.h | 10 +- .../lookup/sort/sort_lookup_store_factory.cpp | 3 +- src/paimon/common/sst/sst_file_io_test.cpp | 13 +- src/paimon/common/sst/sst_file_reader.cpp | 8 +- src/paimon/common/sst/sst_file_reader.h | 9 +- .../global_index/global_index_evaluator.h | 4 +- .../global_index_evaluator_impl.cpp | 63 ++-- .../global_index_evaluator_impl.h | 10 +- .../global_index/global_index_scan_impl.cpp | 30 +- .../global_index/global_index_scan_impl.h | 2 +- .../global_index/global_index_write_task.cpp | 50 ++-- .../row_range_global_index_scanner_impl.cpp | 1 - .../source/data_evolution_batch_scan.cpp | 28 +- .../table/source/data_evolution_batch_scan.h | 2 +- .../lucene/lucene_global_index_reader.cpp | 5 +- .../lucene/lucene_global_index_reader.h | 8 +- .../lucene/lucene_global_index_test.cpp | 5 +- .../lucene/lucene_global_index_writer.cpp | 8 +- .../lucene/lucene_global_index_writer.h | 2 +- .../lumina/lumina_global_index.cpp | 35 +-- .../global_index/lumina/lumina_global_index.h | 40 ++- .../lumina/lumina_global_index_test.cpp | 14 +- test/inte/global_index_test.cpp | 55 ++-- .../btree_test_float_50.bin | 3 + .../btree_test_float_50.bin.meta | 3 + .../btree_test_float_50.csv | 51 ++++ 50 files changed, 1454 insertions(+), 444 deletions(-) create mode 100644 src/paimon/common/global_index/btree/btree_file_meta_selector.cpp create mode 100644 src/paimon/common/global_index/btree/btree_file_meta_selector.h create mode 100644 src/paimon/common/global_index/btree/btree_file_meta_selector_test.cpp create mode 100644 src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp create mode 100644 src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h create mode 100644 src/paimon/common/global_index/global_index_utils.h create mode 100644 test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin create mode 100644 test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin.meta create mode 100644 test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.csv diff --git a/include/paimon/global_index/global_index_io_meta.h b/include/paimon/global_index/global_index_io_meta.h index a8dacd232..e9da91cba 100644 --- a/include/paimon/global_index/global_index_io_meta.h +++ b/include/paimon/global_index/global_index_io_meta.h @@ -25,17 +25,12 @@ namespace paimon { /// Metadata describing a single file entry in a global index. struct PAIMON_EXPORT GlobalIndexIOMeta { - GlobalIndexIOMeta(const std::string& _file_path, int64_t _file_size, int64_t _range_end, + GlobalIndexIOMeta(const std::string& _file_path, int64_t _file_size, const std::shared_ptr& _metadata) - : file_path(_file_path), - file_size(_file_size), - range_end(_range_end), - metadata(_metadata) {} + : file_path(_file_path), file_size(_file_size), metadata(_metadata) {} std::string file_path; int64_t file_size; - /// The inclusive range end covered by this file (i.e., the last local row id). - int64_t range_end; /// Optional binary metadata associated with the file, such as serialized /// secondary index structures or inline index bytes. /// May be null if no additional metadata is available. diff --git a/include/paimon/global_index/global_index_writer.h b/include/paimon/global_index/global_index_writer.h index e7f24a56f..0665698ad 100644 --- a/include/paimon/global_index/global_index_writer.h +++ b/include/paimon/global_index/global_index_writer.h @@ -34,9 +34,10 @@ class PAIMON_EXPORT GlobalIndexWriter { /// /// @param arrow_array A valid C ArrowArray pointer representing a struct array. /// Must not be nullptr, and must conform to the expected schema. + /// @param relative_row_ids local row id calculated by {@code row_id - range.from}. /// @return `Status::OK()` on success; otherwise, an error indicating malformed /// input, I/O failure, or unsupported type, etc. - virtual Status AddBatch(::ArrowArray* arrow_array) = 0; + virtual Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) = 0; /// Finalizes the index build process and returns metadata for persisted index. virtual Result> Finish() = 0; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 4f189a618..d6b3fbc7b 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -416,6 +416,7 @@ if(PAIMON_BUILD_TESTS) common/global_index/btree/key_serializer_test.cpp common/global_index/btree/btree_global_index_integration_test.cpp common/global_index/btree/btree_compatibility_test.cpp + common/global_index/btree/btree_file_meta_selector_test.cpp common/global_index/rangebitmap/range_bitmap_global_index_test.cpp common/global_index/wrap/file_index_reader_wrapper_test.cpp common/io/byte_array_input_stream_test.cpp diff --git a/src/paimon/common/global_index/CMakeLists.txt b/src/paimon/common/global_index/CMakeLists.txt index d61402fbf..c4e7d9548 100644 --- a/src/paimon/common/global_index/CMakeLists.txt +++ b/src/paimon/common/global_index/CMakeLists.txt @@ -20,7 +20,9 @@ set(PAIMON_GLOBAL_INDEX_SRC btree/btree_global_indexer.cpp btree/btree_global_index_reader.cpp btree/btree_global_index_writer.cpp + btree/btree_file_meta_selector.cpp btree/btree_index_meta.cpp + btree/lazy_filtered_btree_reader.cpp btree/key_serializer.cpp rangebitmap/range_bitmap_global_index.cpp rangebitmap/range_bitmap_global_index_factory.cpp) diff --git a/src/paimon/common/global_index/bitmap/bitmap_global_index.cpp b/src/paimon/common/global_index/bitmap/bitmap_global_index.cpp index 6bca15a73..3e8d75346 100644 --- a/src/paimon/common/global_index/bitmap/bitmap_global_index.cpp +++ b/src/paimon/common/global_index/bitmap/bitmap_global_index.cpp @@ -42,9 +42,9 @@ Result> BitmapGlobalIndex::CreateReader( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr reader, index_->CreateReader(arrow_schema, /*start=*/0, meta.file_size, in, pool)); - auto transform = [range_end = meta.range_end](const std::shared_ptr& result) + auto transform = [](const std::shared_ptr& result) -> Result> { - return FileIndexReaderWrapper::ToGlobalIndexResult(range_end, result); + return FileIndexReaderWrapper::ToGlobalIndexResult(result); }; return std::make_shared(reader, transform); } diff --git a/src/paimon/common/global_index/bitmap/bitmap_global_index_test.cpp b/src/paimon/common/global_index/bitmap/bitmap_global_index_test.cpp index d8e88944b..1d18d1347 100644 --- a/src/paimon/common/global_index/bitmap/bitmap_global_index_test.cpp +++ b/src/paimon/common/global_index/bitmap/bitmap_global_index_test.cpp @@ -75,7 +75,9 @@ class BitmapGlobalIndexTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array)); + std::vector row_ids(array->length()); + std::iota(row_ids.begin(), row_ids.end(), 0); + PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array, std::move(row_ids))); PAIMON_ASSIGN_OR_RAISE(auto result_metas, global_writer->Finish()); // check meta if (result_metas.empty()) { @@ -85,7 +87,6 @@ class BitmapGlobalIndexTest : public ::testing::Test { auto file_name = PathUtil::GetName(result_metas[0].file_path); EXPECT_TRUE(StringUtils::StartsWith(file_name, "bitmap-global-index-")); EXPECT_TRUE(StringUtils::EndsWith(file_name, ".index")); - EXPECT_EQ(result_metas[0].range_end, expected_range.to); EXPECT_FALSE(result_metas[0].metadata); return result_metas[0]; } @@ -161,7 +162,7 @@ TEST_F(BitmapGlobalIndexTest, TestStringType) { // greater than return REMAIN file index result, will convert to all range global index // result - CheckResult(reader->VisitGreaterThan(lit_c).value(), {0, 1, 2, 3, 4}); + ASSERT_FALSE(reader->VisitGreaterThan(lit_c).value()); // test visit vector search ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(std::make_shared( @@ -169,11 +170,11 @@ TEST_F(BitmapGlobalIndexTest, TestStringType) { std::nullopt, std::map())), "FileIndexReaderWrapper is not supposed to handle vector search query"); // test VisitStartsWith, VisitEndsWith, VisitContains, VisitLike, VisitFullTextSearch - CheckResult(reader->VisitStartsWith(lit_c).value(), {0, 1, 2, 3, 4}); - CheckResult(reader->VisitEndsWith(lit_c).value(), {0, 1, 2, 3, 4}); - CheckResult(reader->VisitContains(lit_c).value(), {0, 1, 2, 3, 4}); - CheckResult(reader->VisitLike(lit_c).value(), {0, 1, 2, 3, 4}); - CheckResult(reader->VisitFullTextSearch(nullptr).value(), {0, 1, 2, 3, 4}); + ASSERT_FALSE(reader->VisitStartsWith(lit_c).value()); + ASSERT_FALSE(reader->VisitEndsWith(lit_c).value()); + ASSERT_FALSE(reader->VisitContains(lit_c).value()); + ASSERT_FALSE(reader->VisitLike(lit_c).value()); + ASSERT_FALSE(reader->VisitFullTextSearch(nullptr).value()); }; { diff --git a/src/paimon/common/global_index/btree/btree_compatibility_test.cpp b/src/paimon/common/global_index/btree/btree_compatibility_test.cpp index 56a889780..3b58395b8 100644 --- a/src/paimon/common/global_index/btree/btree_compatibility_test.cpp +++ b/src/paimon/common/global_index/btree/btree_compatibility_test.cpp @@ -34,19 +34,19 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { class BTreeCompatibilityTest : public ::testing::Test { - protected: - struct CsvRecord { - int64_t row_id; - std::string key; // "NULL" if is_null - bool is_null; - }; - + public: void SetUp() override { pool_ = GetDefaultPool(); ASSERT_OK_AND_ASSIGN(fs_, FileSystemFactory::Get("local", "/", {})); data_dir_ = GetDataDir() + "/global_index/btree/btree_compatibility_data"; } + struct CsvRecord { + int64_t row_id; + std::string key; // "NULL" if is_null + bool is_null; + }; + std::string ReadFileAsString(const std::string& path) const { EXPECT_OK_AND_ASSIGN(auto input, fs_->Open(path)); EXPECT_OK_AND_ASSIGN(auto length, input->Length()); @@ -110,13 +110,13 @@ class BTreeCompatibilityTest : public ::testing::Test { Result> CreateReaderFromFiles( const std::string& bin_path, const std::string& meta_path, - const std::shared_ptr& arrow_type, int64_t record_count) const { + const std::shared_ptr& arrow_type) const { auto meta_str = ReadFileAsString(meta_path); std::shared_ptr meta_bytes = Bytes::AllocateBytes(meta_str, pool_.get()); PAIMON_ASSIGN_OR_RAISE(auto file_status, fs_->GetFileStatus(bin_path)); auto file_size = static_cast(file_status->GetLen()); - GlobalIndexIOMeta io_meta(bin_path, file_size, record_count - 1, meta_bytes); + GlobalIndexIOMeta io_meta(bin_path, file_size, meta_bytes); std::vector metas = {io_meta}; auto schema = arrow::schema({arrow::field("testField", arrow_type)}); @@ -125,8 +125,8 @@ class BTreeCompatibilityTest : public ::testing::Test { auto file_reader = std::make_shared(fs_); std::map options; - BTreeGlobalIndexer indexer(options); - return indexer.CreateReader(c_schema.get(), file_reader, metas, pool_); + PAIMON_ASSIGN_OR_RAISE(auto indexer, BTreeGlobalIndexer::Create(options)); + return indexer->CreateReader(c_schema.get(), file_reader, metas, pool_); } void RunIntQueries(const std::shared_ptr& reader, @@ -370,6 +370,175 @@ class BTreeCompatibilityTest : public ::testing::Test { } } + // Run float-type queries against a reader with CSV records as ground truth + void RunFloatQueries(const std::shared_ptr& reader, + const std::vector& records) const { + // VisitIsNull + { + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNull()); + auto actual_ids = CollectRowIds(result); + auto expected_ids = + CollectMatchingRows(records, [](const CsvRecord& r) { return r.is_null; }); + ASSERT_EQ(actual_ids, expected_ids); + } + + // VisitIsNotNull + { + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNotNull()); + auto actual_ids = CollectRowIds(result); + auto expected_ids = + CollectMatchingRows(records, [](const CsvRecord& r) { return !r.is_null; }); + ASSERT_EQ(actual_ids, expected_ids); + } + // check special value + { + float value = -INFINITY; + Literal literal(value); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, [](const CsvRecord& r) { return !r.is_null && r.key == "-infinity"; }); + ASSERT_EQ(actual_ids, expected_ids); + } + { + float value = INFINITY; + Literal literal(value); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, [](const CsvRecord& r) { return !r.is_null && r.key == "+infinity"; }); + ASSERT_EQ(actual_ids, expected_ids); + } + { + Literal literal(static_cast(std::nan(""))); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, [](const CsvRecord& r) { return !r.is_null && r.key == "NaN"; }); + ASSERT_EQ(actual_ids, expected_ids); + } + { + Literal literal(static_cast(-0.00f)); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, [](const CsvRecord& r) { return !r.is_null && r.key == "-0.00f"; }); + ASSERT_EQ(actual_ids, expected_ids); + } + { + Literal literal(static_cast(0.00f)); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, [](const CsvRecord& r) { return !r.is_null && r.key == "0.00f"; }); + ASSERT_EQ(actual_ids, expected_ids); + } + + // VisitEqual for the first non-null key + for (const auto& rec : records) { + if (!rec.is_null && rec.key != "-infinity") { + float key_val = std::stof(rec.key); + Literal literal(key_val); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows(records, [key_val](const CsvRecord& r) { + return !r.is_null && std::stof(r.key) == key_val; + }); + ASSERT_EQ(actual_ids, expected_ids); + break; + } + } + + // VisitEqual for the last non-null key + for (auto it = records.rbegin(); it != records.rend(); ++it) { + if (!it->is_null && it->key != "+infinity" && it->key != "NaN") { + float key_val = std::stof(it->key); + Literal literal(key_val); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows(records, [key_val](const CsvRecord& r) { + return !r.is_null && std::stof(r.key) == key_val; + }); + ASSERT_EQ(actual_ids, expected_ids); + break; + } + } + + // VisitEqual for a non-existent key + { + Literal literal(static_cast(-99.99)); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal)); + auto actual_ids = CollectRowIds(result); + ASSERT_TRUE(actual_ids.empty()); + } + + float mid_key = -1; + { + std::vector non_null_keys; + for (const auto& rec : records) { + if (!rec.is_null) { + non_null_keys.push_back(std::stof(rec.key)); + } + } + ASSERT_FALSE(non_null_keys.empty()); + mid_key = non_null_keys[non_null_keys.size() / 2]; + } + // VisitLessThan for a mid-range key + { + Literal literal(mid_key); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessThan(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows(records, [mid_key](const CsvRecord& r) { + return !r.is_null && std::stof(r.key) < mid_key; + }); + ASSERT_EQ(actual_ids, expected_ids); + } + + // VisitGreaterOrEqual for a mid-range key + { + Literal literal(mid_key); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterOrEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows(records, [mid_key](const CsvRecord& r) { + return !r.is_null && (std::stof(r.key) >= mid_key || r.key == "NaN"); + }); + ASSERT_EQ(actual_ids, expected_ids); + } + + // VisitIn for multiple keys + { + float k1 = -9.27f; + float k2 = 62.91f; + float k3 = 108.17f; + std::vector in_literals = {Literal(k1), Literal(k2), Literal(k3)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(in_literals)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows(records, [k1, k2, k3](const CsvRecord& r) { + if (r.is_null) { + return false; + } + float v = std::stof(r.key); + return v == k1 || v == k2 || v == k3; + }); + ASSERT_EQ(actual_ids, expected_ids); + } + + // VisitNotEqual for the first non-null key + for (const auto& rec : records) { + if (!rec.is_null) { + float key_val = std::stof(rec.key); + Literal literal(key_val); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotEqual(literal)); + auto actual_ids = CollectRowIds(result); + auto expected_ids = CollectMatchingRows( + records, + [key_val, rec](const CsvRecord& r) { return !r.is_null && r.key != rec.key; }); + ASSERT_EQ(actual_ids, expected_ids); + break; + } + } + } + class LocalGlobalIndexFileReader : public GlobalIndexFileReader { public: explicit LocalGlobalIndexFileReader(const std::shared_ptr& fs) : fs_(fs) {} @@ -383,6 +552,7 @@ class BTreeCompatibilityTest : public ::testing::Test { std::shared_ptr fs_; }; + private: std::shared_ptr pool_; std::shared_ptr fs_; std::string data_dir_; @@ -399,7 +569,7 @@ TEST_F(BTreeCompatibilityTest, ReadAndQueryIntData) { ASSERT_EQ(static_cast(records.size()), count); ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow::int32(), count)); + CreateReaderFromFiles(bin_path, meta_path, arrow::int32())); RunIntQueries(reader, records); } } @@ -415,11 +585,25 @@ TEST_F(BTreeCompatibilityTest, ReadAndQueryStringData) { ASSERT_EQ(static_cast(records.size()), count); ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow::utf8(), count)); + CreateReaderFromFiles(bin_path, meta_path, arrow::utf8())); RunStringQueries(reader, records); } } +TEST_F(BTreeCompatibilityTest, ReadAndQueryFloatData) { + int32_t count = 50; + std::string prefix = "btree_test_float_" + std::to_string(count); + std::string bin_path = data_dir_ + "/" + prefix + ".bin"; + std::string meta_path = bin_path + ".meta"; + std::string csv_path = data_dir_ + "/" + prefix + ".csv"; + + auto records = ParseCsvFile(csv_path); + ASSERT_EQ(static_cast(records.size()), count); + + ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow::float32())); + RunFloatQueries(reader, records); +} + TEST_F(BTreeCompatibilityTest, AllNulls) { std::string prefix = "btree_test_int_all_nulls"; std::string bin_path = data_dir_ + "/" + prefix + ".bin"; @@ -430,8 +614,7 @@ TEST_F(BTreeCompatibilityTest, AllNulls) { ASSERT_FALSE(records.empty()); auto count = static_cast(records.size()); - ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow::int32(), count)); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow::int32())); // All rows should be null { @@ -469,8 +652,7 @@ TEST_F(BTreeCompatibilityTest, NoNulls) { ASSERT_FALSE(records.empty()); auto count = static_cast(records.size()); - ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow::int32(), count)); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow::int32())); // No rows should be null { @@ -538,8 +720,7 @@ TEST_F(BTreeCompatibilityTest, DuplicateKeys) { ASSERT_FALSE(records.empty()); auto count = static_cast(records.size()); - ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow::int32(), count)); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow::int32())); // VisitIsNull { @@ -703,8 +884,7 @@ TEST_F(BTreeCompatibilityTest, RowCountConsistency) { ASSERT_FALSE(records.empty()); auto count = static_cast(records.size()); - ASSERT_OK_AND_ASSIGN(auto reader, - CreateReaderFromFiles(bin_path, meta_path, arrow_type, count)); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow_type)); ASSERT_OK_AND_ASSIGN(auto null_result, reader->VisitIsNull()); auto null_ids = CollectRowIds(null_result); 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 000000000..192bab7a8 --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector.cpp @@ -0,0 +1,172 @@ +/* + * 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/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) { + 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) { + return Filter([this, &literal](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal min_key, DeserializeKey(meta.FirstKey())); + PAIMON_ASSIGN_OR_RAISE(Literal max_key, DeserializeKey(meta.LastKey())); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_min, literal.CompareTo(min_key)); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_max, literal.CompareTo(max_key)); + 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 + return Filter([this, &literal](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal min_key, DeserializeKey(meta.FirstKey())); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, min_key.CompareTo(literal)); + return cmp < 0; + }); +} + +Result> BTreeFileMetaSelector::VisitLessOrEqual( + const Literal& literal) { + // file.minKey <= literal + return Filter([this, &literal](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal min_key, DeserializeKey(meta.FirstKey())); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, min_key.CompareTo(literal)); + return cmp <= 0; + }); +} + +Result> BTreeFileMetaSelector::VisitGreaterThan( + const Literal& literal) { + // file.maxKey > literal + return Filter([this, &literal](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal max_key, DeserializeKey(meta.LastKey())); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, max_key.CompareTo(literal)); + return cmp > 0; + }); +} + +Result> BTreeFileMetaSelector::VisitGreaterOrEqual( + const Literal& literal) { + // file.maxKey >= literal + return Filter([this, &literal](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal max_key, DeserializeKey(meta.LastKey())); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp, max_key.CompareTo(literal)); + return cmp >= 0; + }); +} + +Result> BTreeFileMetaSelector::VisitIn( + const std::vector& literals) { + return Filter([this, &literals](const BTreeIndexMeta& meta) -> Result { + if (meta.OnlyNulls()) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(Literal min_key, DeserializeKey(meta.FirstKey())); + PAIMON_ASSIGN_OR_RAISE(Literal max_key, DeserializeKey(meta.LastKey())); + for (const auto& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_min, literal.CompareTo(min_key)); + PAIMON_ASSIGN_OR_RAISE(int32_t cmp_max, literal.CompareTo(max_key)); + 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; +} + +Result BTreeFileMetaSelector::DeserializeKey( + const std::shared_ptr& key_bytes) const { + MemorySlice slice = MemorySlice::Wrap(key_bytes); + return KeySerializer::DeserializeKey(slice, key_type_, pool_.get()); +} + +} // 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 000000000..baeaf19eb --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector.h @@ -0,0 +1,65 @@ +/* + * 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 +#include +#include + +#include "paimon/common/global_index/btree/btree_index_meta.h" +#include "paimon/common/global_index/btree/key_serializer.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 DeserializeKey(const std::shared_ptr& key_bytes) const; + + std::vector>> files_; + std::shared_ptr key_type_; + std::shared_ptr pool_; +}; + +} // 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 000000000..c24c3c61b --- /dev/null +++ b/src/paimon/common/global_index/btree/btree_file_meta_selector_test.cpp @@ -0,0 +1,228 @@ +/* + * 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/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, [19,25]=file5 + // 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_global_index_factory.cpp b/src/paimon/common/global_index/btree/btree_global_index_factory.cpp index f07eeb85d..199b2d706 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_factory.cpp +++ b/src/paimon/common/global_index/btree/btree_global_index_factory.cpp @@ -25,7 +25,7 @@ const char BTreeGlobalIndexerFactory::IDENTIFIER[] = "btree-global"; Result> BTreeGlobalIndexerFactory::Create( const std::map& options) const { - return std::make_unique(options); + return BTreeGlobalIndexer::Create(options); } REGISTER_PAIMON_FACTORY(BTreeGlobalIndexerFactory); diff --git a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp index 15f7bf5bd..f015988f4 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp +++ b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp @@ -124,7 +124,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadIntData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "4096"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -151,7 +151,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -278,7 +278,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadStringData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("str_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -304,7 +304,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadStringData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -446,7 +446,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBigIntData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("bigint_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -469,7 +469,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBigIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -538,7 +538,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadFloatData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("float_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -561,7 +561,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadFloatData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -629,7 +629,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDoubleData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("double_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -652,7 +652,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDoubleData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -720,7 +720,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNonNull) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -742,7 +742,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNonNull) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -829,7 +829,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBoolData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("bool_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -853,7 +853,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBoolData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -906,7 +906,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTinyIntData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("tinyint_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -929,7 +929,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTinyIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -988,7 +988,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadSmallIntData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("smallint_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1011,7 +1011,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadSmallIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1071,7 +1071,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampCompactData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("ts_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1094,7 +1094,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1154,7 +1154,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampNonCompactData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("ts_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1178,7 +1178,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampNonCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1242,7 +1242,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalCompactData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("decimal_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1266,7 +1266,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1326,7 +1326,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalNonCompactData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("decimal_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1351,7 +1351,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalNonCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1410,7 +1410,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNull) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1429,7 +1429,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNull) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1498,7 +1498,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadLargeDataWithSmallBlocks) { {BtreeDefs::kBtreeIndexBlockSize, "256"}, {BtreeDefs::kBtreeIndexCacheSize, "1024"}, }; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1544,7 +1544,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadLargeDataWithSmallBlocks) { ArrowArray c_array; ASSERT_TRUE(arrow::ExportArray(*struct_array, &c_array).ok()); - ASSERT_OK(btree_writer->AddBatch(&c_array, batch_row_ids)); + ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(batch_row_ids))); } ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); @@ -1643,7 +1643,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, CreateWriterWithNonStructSchema) { auto file_writer = std::make_shared(fs_, base_path_); std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); // Export a plain int32 type (not struct) as ArrowSchema auto plain_type = arrow::int32(); @@ -1661,11 +1661,11 @@ TEST_P(BTreeGlobalIndexIntegrationTest, CreateReaderWithMultipleMetas) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); // Provide two fake metas - GlobalIndexIOMeta meta1("fake_path_1", 100, 10, nullptr); - GlobalIndexIOMeta meta2("fake_path_2", 200, 20, nullptr); + GlobalIndexIOMeta meta1("fake_path_1", 100, nullptr); + GlobalIndexIOMeta meta2("fake_path_2", 200, nullptr); std::vector metas = {meta1, meta2}; ASSERT_NOK_WITH_MSG(indexer->CreateReader(c_schema.get(), file_reader, metas, pool_), @@ -1683,9 +1683,9 @@ TEST_P(BTreeGlobalIndexIntegrationTest, CreateReaderWithMultiFieldSchema) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); - GlobalIndexIOMeta meta("fake_path", 100, 10, nullptr); + GlobalIndexIOMeta meta("fake_path", 100, nullptr); std::vector metas = {meta}; ASSERT_NOK_WITH_MSG(indexer->CreateReader(c_schema.get(), file_reader, metas, pool_), @@ -1710,14 +1710,15 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNullArray) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); ASSERT_TRUE(btree_writer); std::vector row_ids = {0, 1, 2}; - ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(nullptr, row_ids), "ArrowArray is null"); + ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(nullptr, std::move(row_ids)), + "CheckRelativeRowIds failed: null c_arrow_array"); } TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithMismatchedRowIds) { @@ -1727,7 +1728,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithMismatchedRowIds) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1745,8 +1746,9 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithMismatchedRowIds) { // Provide wrong number of row_ids (2 instead of 3) std::vector row_ids = {0, 1}; - ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(&c_array, row_ids), - "row_ids length 2 mismatch arrow_array length 3 when AddBatch"); + ASSERT_NOK_WITH_MSG( + btree_writer->AddBatch(&c_array, std::move(row_ids)), + "relative_row_ids length 2 mismatch arrow_array length 3 in CheckRelativeRowIds"); } TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNonMonotonicKeys) { @@ -1756,7 +1758,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNonMonotonicKeys) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1773,7 +1775,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNonMonotonicKeys) { ArrowArray c_array; ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids = {0, 1, 2}; - ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(&c_array, row_ids), + ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(&c_array, std::move(row_ids)), "Users must keep written keys monotonically incremental"); } @@ -1784,7 +1786,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, FinishWithEmptyData) { std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); auto btree_writer = std::dynamic_pointer_cast(writer); @@ -1812,7 +1814,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, TestIOException) { std::map options = { {BtreeDefs::kBtreeIndexBlockSize, "128"}, {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - auto indexer = std::make_shared(options); + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); // write auto writer_result = indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_); @@ -1828,7 +1830,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, TestIOException) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids = {0, 1, 2, 3, 4, 5}; - CHECK_HOOK_STATUS(btree_writer->AddBatch(&c_array, row_ids), i); + CHECK_HOOK_STATUS(btree_writer->AddBatch(&c_array, std::move(row_ids)), i); auto finish_result = writer->Finish(); CHECK_HOOK_STATUS(finish_result.status(), i); auto metas = std::move(finish_result).value(); diff --git a/src/paimon/common/global_index/btree/btree_global_index_writer.cpp b/src/paimon/common/global_index/btree/btree_global_index_writer.cpp index ea7b5baf4..5eabae3fb 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_writer.cpp +++ b/src/paimon/common/global_index/btree/btree_global_index_writer.cpp @@ -19,10 +19,12 @@ #include #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "fmt/format.h" #include "paimon/common/compression/block_compression_factory.h" #include "paimon/common/global_index/btree/btree_defs.h" #include "paimon/common/global_index/btree/key_serializer.h" +#include "paimon/common/global_index/global_index_utils.h" #include "paimon/common/memory/memory_slice_output.h" #include "paimon/common/predicate/literal_converter.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -66,18 +68,11 @@ BTreeGlobalIndexWriter::BTreeGlobalIndexWriter( sst_writer_(std::move(sst_writer)) {} Status BTreeGlobalIndexWriter::AddBatch(::ArrowArray* arrow_array, - const std::vector& row_ids) { - if (!arrow_array) { - return Status::Invalid("ArrowArray is null"); - } + std::vector&& relative_row_ids) { + PAIMON_RETURN_NOT_OK(GlobalIndexUtils::CheckRelativeRowIds( + arrow_array, relative_row_ids, /*expected_next_row_id=*/std::nullopt)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(arrow_array, arrow_type_)); - if (static_cast(row_ids.size()) != array->length()) { - return Status::Invalid( - fmt::format("row_ids length {} mismatch arrow_array length {} when AddBatch to " - "BTreeGlobalIndexWriter", - row_ids.size(), array->length())); - } auto struct_array = std::dynamic_pointer_cast(array); PAIMON_RETURN_NOT_OK(Preconditions::CheckNotNull( struct_array, "arrow array must be struct array when AddBatch to BTreeGlobalIndexWriter")); @@ -92,9 +87,8 @@ Status BTreeGlobalIndexWriter::AddBatch(::ArrowArray* arrow_array, LiteralConverter::ConvertLiteralsFromArray(*value_array, /*own_data=*/true)); for (size_t i = 0; i < literals.size(); ++i) { - int64_t row_id = row_ids[i]; + int64_t row_id = relative_row_ids[i]; const auto& literal = literals[i]; - max_row_id_ = std::max(max_row_id_, row_id); if (literal.IsNull()) { // Track null values null_bitmap_.Add(row_id); @@ -202,7 +196,7 @@ Result> BTreeGlobalIndexWriter::Finish() { // Create GlobalIndexIOMeta std::string file_path = file_writer_->ToPath(index_file_name_); PAIMON_ASSIGN_OR_RAISE(int64_t file_size, file_writer_->GetFileSize(index_file_name_)); - GlobalIndexIOMeta io_meta(file_path, file_size, max_row_id_, meta_bytes); + GlobalIndexIOMeta io_meta(file_path, file_size, meta_bytes); return std::vector{io_meta}; } diff --git a/src/paimon/common/global_index/btree/btree_global_index_writer.h b/src/paimon/common/global_index/btree/btree_global_index_writer.h index de72773d7..a7882fda7 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_writer.h +++ b/src/paimon/common/global_index/btree/btree_global_index_writer.h @@ -65,12 +65,7 @@ class BTreeGlobalIndexWriter : public GlobalIndexWriter { ~BTreeGlobalIndexWriter() override = default; - Status AddBatch(::ArrowArray* arrow_array) override { - // TODO(xinyu.lxy): refactor AddBatch with relative row ids - return Status::Invalid("BTreeGlobalIndexWriter not support AddBatch without row_ids"); - } - - Status AddBatch(::ArrowArray* arrow_array, const std::vector& row_ids); + Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) override; /// Finish writing and return the index metadata. Result> Finish() override; @@ -100,8 +95,6 @@ class BTreeGlobalIndexWriter : public GlobalIndexWriter { std::shared_ptr output_stream_; std::unique_ptr sst_writer_; - // TODO(xinyu.lxy): remove it when GlobalIndexIOMeta is updated - int64_t max_row_id_ = -1; std::optional first_key_; std::optional last_key_; diff --git a/src/paimon/common/global_index/btree/btree_global_indexer.cpp b/src/paimon/common/global_index/btree/btree_global_indexer.cpp index 9e96acc23..43c4bfb49 100644 --- a/src/paimon/common/global_index/btree/btree_global_indexer.cpp +++ b/src/paimon/common/global_index/btree/btree_global_indexer.cpp @@ -21,10 +21,10 @@ #include "arrow/c/bridge.h" #include "paimon/common/compression/block_compression_factory.h" #include "paimon/common/global_index/btree/btree_file_footer.h" -#include "paimon/common/global_index/btree/btree_global_index_reader.h" #include "paimon/common/global_index/btree/btree_global_index_writer.h" #include "paimon/common/global_index/btree/btree_index_meta.h" #include "paimon/common/global_index/btree/key_serializer.h" +#include "paimon/common/global_index/btree/lazy_filtered_btree_reader.h" #include "paimon/common/memory/memory_slice.h" #include "paimon/common/memory/memory_slice_input.h" #include "paimon/common/options/memory_size.h" @@ -37,6 +37,22 @@ #include "paimon/memory/bytes.h" #include "paimon/utils/roaring_bitmap64.h" namespace paimon { +Result> BTreeGlobalIndexer::Create( + const std::map& options) { + // parse cache options + PAIMON_ASSIGN_OR_RAISE(std::string cache_size_str, OptionsUtils::GetValueFromMap( + options, BtreeDefs::kBtreeIndexCacheSize, + BtreeDefs::kDefaultBtreeIndexCacheSize)); + PAIMON_ASSIGN_OR_RAISE(int64_t cache_size, MemorySize::ParseBytes(cache_size_str)); + + PAIMON_ASSIGN_OR_RAISE( + double high_priority_pool_ratio, + OptionsUtils::GetValueFromMap(options, BtreeDefs::kBtreeIndexHighPriorityPoolRatio, + BtreeDefs::kDefaultBtreeIndexHighPriorityPoolRatio)); + auto cache_manager = std::make_shared(cache_size, high_priority_pool_ratio); + return std::unique_ptr(new BTreeGlobalIndexer(cache_manager, options)); +} + Result> BTreeGlobalIndexer::CreateWriter( const std::string& field_name, ::ArrowSchema* arrow_schema, const std::shared_ptr& file_writer, @@ -75,113 +91,13 @@ Result> BTreeGlobalIndexer::CreateReader( // Get field type from arrow schema PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, arrow::ImportSchema(arrow_schema)); - if (files.size() != 1) { - return Status::Invalid( - "invalid GlobalIndexIOMeta for BTreeGlobalIndex, exist multiple metas"); - } - const auto& meta = files[0]; if (schema->num_fields() != 1) { return Status::Invalid( "invalid schema for BTreeGlobalIndexReader, supposed to have single field."); } auto key_type = schema->field(0)->type(); - // Create comparator based on field type - auto comparator = KeySerializer::CreateComparator(key_type, pool); - // get min, max key from meta data - auto index_meta = BTreeIndexMeta::Deserialize(meta.metadata, pool.get()); - std::optional min_key; - std::optional max_key; - if (index_meta->FirstKey()) { - PAIMON_ASSIGN_OR_RAISE( - min_key, KeySerializer::DeserializeKey(MemorySlice::Wrap(index_meta->FirstKey()), - key_type, pool.get())); - } - if (index_meta->LastKey()) { - PAIMON_ASSIGN_OR_RAISE( - max_key, KeySerializer::DeserializeKey(MemorySlice::Wrap(index_meta->LastKey()), - key_type, pool.get())); - } - - // parse read options - PAIMON_ASSIGN_OR_RAISE( - std::string cache_size_str, - OptionsUtils::GetValueFromMap(options_, BtreeDefs::kBtreeIndexCacheSize, - BtreeDefs::kDefaultBtreeIndexCacheSize)); - PAIMON_ASSIGN_OR_RAISE(int64_t cache_size, MemorySize::ParseBytes(cache_size_str)); - - PAIMON_ASSIGN_OR_RAISE( - double high_priority_pool_ratio, - OptionsUtils::GetValueFromMap(options_, BtreeDefs::kBtreeIndexHighPriorityPoolRatio, - BtreeDefs::kDefaultBtreeIndexHighPriorityPoolRatio)); - // TODO(xinyu.lxy): pass cache_manager from param. - auto cache_manager = std::make_shared(cache_size, high_priority_pool_ratio); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, - file_reader->GetInputStream(meta.file_path)); - auto block_cache = std::make_shared(meta.file_path, in, cache_manager, pool); - // read footer - PAIMON_ASSIGN_OR_RAISE(MemorySegment segment, - block_cache->GetBlock(meta.file_size - BTreeFileFooter::kEncodingLength, - BTreeFileFooter::kEncodingLength, true, - /*decompress_func=*/nullptr)); - auto footer_slice = MemorySlice::Wrap(segment); - auto footer_input = footer_slice.ToInput(); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr footer, - BTreeFileFooter::Read(&footer_input)); - // prepare null_bitmap - PAIMON_ASSIGN_OR_RAISE(RoaringBitmap64 null_bitmap, - ReadNullBitmap(block_cache, footer->GetNullBitmapHandle(), pool.get())); - - // Close the temporary block_cache to remove its entries from the shared LRU cache. - // This prevents use-after-free: the eviction callback captures `this` (the BlockCache), - // and if the BlockCache is destroyed without closing, a later eviction would invoke - // the callback on a dangling pointer. - block_cache->Close(); - - // create SST file reader with footer information - // TODO(xinyu.lxy): pass block cache to SstFileReader rather than cache_manager - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr sst_file_reader, - SstFileReader::Create(in, footer->GetIndexBlockHandle(), footer->GetBloomFilterHandle(), - comparator, cache_manager, pool)); - - return std::make_shared(sst_file_reader, std::move(null_bitmap), - min_key, max_key, key_type, pool); -} - -Result BTreeGlobalIndexer::ReadNullBitmap( - const std::shared_ptr& cache, const std::optional& block_handle, - MemoryPool* pool) { - RoaringBitmap64 null_bitmap; - if (!block_handle.has_value()) { - return null_bitmap; - } - - // read bytes and crc value - PAIMON_ASSIGN_OR_RAISE( - MemorySegment segment, - cache->GetBlock(block_handle->Offset(), block_handle->Size() + 4, /*is_index=*/false, - /*decompress_func=*/nullptr)); - - auto slice = MemorySlice::Wrap(segment); - auto slice_input = slice.ToInput(); - - // read null bitmap data - auto null_bitmap_bytes = slice_input.ReadSlice(block_handle->Size()).CopyBytes(pool); - // Calculate CRC32C checksum - uint32_t crc_value = CRC32C::calculate(null_bitmap_bytes->data(), null_bitmap_bytes->size()); - int32_t expected_crc_value = slice_input.ReadInt(); - - // Verify CRC checksum - if (crc_value != static_cast(expected_crc_value)) { - return Status::Invalid(fmt::format( - "CRC check failure during decoding null bitmap. Expected: {}, Calculated: {}", - expected_crc_value, crc_value)); - } - - // deserialize null bitmap - PAIMON_RETURN_NOT_OK( - null_bitmap.Deserialize(null_bitmap_bytes->data(), null_bitmap_bytes->size())); - return null_bitmap; + return std::make_shared(files, key_type, file_reader, cache_manager_, + pool); } } // namespace paimon diff --git a/src/paimon/common/global_index/btree/btree_global_indexer.h b/src/paimon/common/global_index/btree/btree_global_indexer.h index e1eda2478..bf05e704d 100644 --- a/src/paimon/common/global_index/btree/btree_global_indexer.h +++ b/src/paimon/common/global_index/btree/btree_global_indexer.h @@ -50,8 +50,8 @@ namespace paimon { class BTreeGlobalIndexer : public GlobalIndexer { public: - explicit BTreeGlobalIndexer(const std::map& options) - : options_(options) {} + static Result> Create( + const std::map& options); Result> CreateWriter( const std::string& field_name, ::ArrowSchema* arrow_schema, @@ -64,11 +64,12 @@ class BTreeGlobalIndexer : public GlobalIndexer { const std::shared_ptr& pool) const override; private: - static Result ReadNullBitmap(const std::shared_ptr& cache, - const std::optional& block_handle, - MemoryPool* pool); + BTreeGlobalIndexer(const std::shared_ptr& cache_manager, + const std::map& options) + : cache_manager_(cache_manager), options_(options) {} private: + std::shared_ptr cache_manager_; std::map options_; }; diff --git a/src/paimon/common/global_index/btree/key_serializer.cpp b/src/paimon/common/global_index/btree/key_serializer.cpp index 1aab2951a..dc5456907 100644 --- a/src/paimon/common/global_index/btree/key_serializer.cpp +++ b/src/paimon/common/global_index/btree/key_serializer.cpp @@ -67,7 +67,6 @@ Result> KeySerializer::SerializeKey( case FieldType::FLOAT: { MemorySliceOutput output(4, pool); output.Reset(); - // TODO(xinyu): check java floatToIntBits auto fvalue = literal.GetValue(); int32_t ivalue; memcpy(&ivalue, &fvalue, sizeof(float)); @@ -77,7 +76,6 @@ Result> KeySerializer::SerializeKey( case FieldType::DOUBLE: { MemorySliceOutput output(8, pool); output.Reset(); - // TODO(xinyu): check java doubleToLongBits auto dvalue = literal.GetValue(); int64_t ivalue; memcpy(&ivalue, &dvalue, sizeof(double)); diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp new file mode 100644 index 000000000..f7dc0efec --- /dev/null +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp @@ -0,0 +1,270 @@ +/* + * 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/global_index/btree/lazy_filtered_btree_reader.h" + +#include + +#include "paimon/common/global_index/btree/btree_file_footer.h" +#include "paimon/common/global_index/btree/btree_global_index_reader.h" +#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/common/memory/memory_slice_input.h" +#include "paimon/common/sst/block_cache.h" +#include "paimon/common/sst/sst_file_reader.h" +#include "paimon/common/utils/crc32c.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { +LazyFilteredBTreeReader::LazyFilteredBTreeReader( + const std::vector& files, const std::shared_ptr& key_type, + const std::shared_ptr& file_reader, + const std::shared_ptr& cache_manager, const std::shared_ptr& pool) + : pool_(pool), + file_selector_(files, key_type, pool), + key_type_(key_type), + file_reader_(file_reader), + cache_manager_(cache_manager) {} + +Result> LazyFilteredBTreeReader::VisitIsNotNull() { + return DispatchVisit([this]() { return file_selector_.VisitIsNotNull(); }, + [](GlobalIndexReader& reader) { return reader.VisitIsNotNull(); }); +} + +Result> LazyFilteredBTreeReader::VisitIsNull() { + return DispatchVisit([this]() { return file_selector_.VisitIsNull(); }, + [](GlobalIndexReader& reader) { return reader.VisitIsNull(); }); +} + +Result> LazyFilteredBTreeReader::VisitEqual( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitEqual(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitEqual(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitNotEqual( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitNotEqual(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitNotEqual(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitLessThan( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitLessThan(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitLessThan(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitLessOrEqual( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitLessOrEqual(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitLessOrEqual(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitGreaterThan( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitGreaterThan(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitGreaterThan(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitGreaterOrEqual( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitGreaterOrEqual(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitGreaterOrEqual(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitIn( + const std::vector& literals) { + return DispatchVisit( + [this, &literals]() { return file_selector_.VisitIn(literals); }, + [&literals](GlobalIndexReader& reader) { return reader.VisitIn(literals); }); +} + +Result> LazyFilteredBTreeReader::VisitNotIn( + const std::vector& literals) { + return DispatchVisit( + [this, &literals]() { return file_selector_.VisitNotIn(literals); }, + [&literals](GlobalIndexReader& reader) { return reader.VisitNotIn(literals); }); +} + +Result> LazyFilteredBTreeReader::VisitStartsWith( + const Literal& prefix) { + return DispatchVisit( + [this, &prefix]() { return file_selector_.VisitStartsWith(prefix); }, + [&prefix](GlobalIndexReader& reader) { return reader.VisitStartsWith(prefix); }); +} + +Result> LazyFilteredBTreeReader::VisitEndsWith( + const Literal& suffix) { + return DispatchVisit( + [this, &suffix]() { return file_selector_.VisitEndsWith(suffix); }, + [&suffix](GlobalIndexReader& reader) { return reader.VisitEndsWith(suffix); }); +} + +Result> LazyFilteredBTreeReader::VisitContains( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitContains(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitContains(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitLike( + const Literal& literal) { + return DispatchVisit( + [this, &literal]() { return file_selector_.VisitLike(literal); }, + [&literal](GlobalIndexReader& reader) { return reader.VisitLike(literal); }); +} + +Result> LazyFilteredBTreeReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + return Status::Invalid("LazyFilteredBTreeReader does not support vector search"); +} + +Result> LazyFilteredBTreeReader::VisitFullTextSearch( + const std::shared_ptr& full_text_search) { + return Status::Invalid("LazyFilteredBTreeReader does not support full text search"); +} + +Result> LazyFilteredBTreeReader::DispatchVisit( + SelectAction select_files, ReaderAction action) { + PAIMON_ASSIGN_OR_RAISE(auto selected_files, select_files()); + if (selected_files.empty()) { + return std::make_shared([]() { return RoaringBitmap64(); }); + } + + std::shared_ptr merged_result = nullptr; + for (const auto& meta : selected_files) { + PAIMON_ASSIGN_OR_RAISE(auto reader, GetOrCreateReader(meta)); + PAIMON_ASSIGN_OR_RAISE(auto result, action(*reader)); + if (result == nullptr) { + continue; + } + if (merged_result == nullptr) { + merged_result = result; + } else { + PAIMON_ASSIGN_OR_RAISE(merged_result, merged_result->Or(result)); + } + } + + if (merged_result == nullptr) { + return Status::Invalid("DispatchVisit cannot return empty result"); + } + return merged_result; +} + +Result> LazyFilteredBTreeReader::GetOrCreateReader( + const GlobalIndexIOMeta& meta) { + auto iterator = reader_cache_.find(meta.file_path); + if (iterator != reader_cache_.end()) { + return iterator->second; + } + PAIMON_ASSIGN_OR_RAISE(auto reader, CreateSingleReader(meta)); + reader_cache_[meta.file_path] = reader; + return reader; +} + +Result> LazyFilteredBTreeReader::CreateSingleReader( + const GlobalIndexIOMeta& meta) { + auto comparator = KeySerializer::CreateComparator(key_type_, pool_); + + // Deserialize min/max keys from meta + auto index_meta = BTreeIndexMeta::Deserialize(meta.metadata, pool_.get()); + std::optional min_key; + std::optional max_key; + if (index_meta->FirstKey()) { + PAIMON_ASSIGN_OR_RAISE( + min_key, KeySerializer::DeserializeKey(MemorySlice::Wrap(index_meta->FirstKey()), + key_type_, pool_.get())); + } + if (index_meta->LastKey()) { + PAIMON_ASSIGN_OR_RAISE( + max_key, KeySerializer::DeserializeKey(MemorySlice::Wrap(index_meta->LastKey()), + key_type_, pool_.get())); + } + + // Open input stream and create block cache + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, + file_reader_->GetInputStream(meta.file_path)); + auto block_cache = + std::make_shared(meta.file_path, input_stream, cache_manager_, pool_); + + // Read footer + PAIMON_ASSIGN_OR_RAISE(MemorySegment footer_segment, + block_cache->GetBlock(meta.file_size - BTreeFileFooter::kEncodingLength, + BTreeFileFooter::kEncodingLength, true, + /*decompress_func=*/nullptr)); + auto footer_slice = MemorySlice::Wrap(footer_segment); + auto footer_input = footer_slice.ToInput(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr footer, + BTreeFileFooter::Read(&footer_input)); + + // Read null bitmap + PAIMON_ASSIGN_OR_RAISE(RoaringBitmap64 null_bitmap, + ReadNullBitmap(block_cache, footer->GetNullBitmapHandle())); + + // Create SST file reader + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr sst_file_reader, + SstFileReader::Create(input_stream, footer->GetIndexBlockHandle(), + footer->GetBloomFilterHandle(), comparator, block_cache, pool_)); + + return std::make_shared(sst_file_reader, std::move(null_bitmap), + min_key, max_key, key_type_, pool_); +} + +Result LazyFilteredBTreeReader::ReadNullBitmap( + const std::shared_ptr& cache, const std::optional& block_handle) { + RoaringBitmap64 null_bitmap; + if (!block_handle.has_value()) { + return null_bitmap; + } + + // Read bytes and CRC value + PAIMON_ASSIGN_OR_RAISE( + MemorySegment segment, + cache->GetBlock(block_handle->Offset(), block_handle->Size() + 4, /*is_index=*/false, + /*decompress_func=*/nullptr)); + + auto slice = MemorySlice::Wrap(segment); + auto slice_input = slice.ToInput(); + + // Read null bitmap data + auto null_bitmap_bytes = slice_input.ReadSlice(block_handle->Size()).CopyBytes(pool_.get()); + + // Calculate and verify CRC32C checksum + uint32_t calculated_crc = + CRC32C::calculate(null_bitmap_bytes->data(), null_bitmap_bytes->size()); + int32_t expected_crc = slice_input.ReadInt(); + if (calculated_crc != static_cast(expected_crc)) { + return Status::Invalid(fmt::format( + "CRC check failure during decoding null bitmap. Expected: {}, Calculated: {}", + expected_crc, calculated_crc)); + } + + // Deserialize null bitmap + PAIMON_RETURN_NOT_OK( + null_bitmap.Deserialize(null_bitmap_bytes->data(), null_bitmap_bytes->size())); + return null_bitmap; +} + +} // namespace paimon diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h new file mode 100644 index 000000000..bd313a692 --- /dev/null +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h @@ -0,0 +1,98 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "paimon/common/global_index/btree/btree_defs.h" +#include "paimon/common/global_index/btree/btree_file_meta_selector.h" +#include "paimon/common/io/cache/cache_manager.h" +#include "paimon/common/sst/block_cache.h" +#include "paimon/common/sst/block_handle.h" +#include "paimon/global_index/global_index_reader.h" +#include "paimon/global_index/io/global_index_file_reader.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { +class LazyFilteredBTreeReader : public GlobalIndexReader { + public: + LazyFilteredBTreeReader(const std::vector& files, + const std::shared_ptr& key_type, + const std::shared_ptr& file_reader, + const std::shared_ptr& cache_manager, + 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; + + Result> VisitVectorSearch( + const std::shared_ptr& vector_search) override; + + Result> VisitFullTextSearch( + const std::shared_ptr& full_text_search) override; + + bool IsThreadSafe() const override { + return false; + } + + std::string GetIndexType() const override { + return BtreeDefs::kIdentifier; + } + + private: + using SelectAction = std::function>()>; + using ReaderAction = + std::function>(GlobalIndexReader&)>; + + Result> DispatchVisit(SelectAction select_files, + ReaderAction action); + Result> GetOrCreateReader(const GlobalIndexIOMeta& meta); + Result> CreateSingleReader(const GlobalIndexIOMeta& meta); + Result ReadNullBitmap(const std::shared_ptr& cache, + const std::optional& block_handle); + + private: + // TODO(lisizhuo.lsz): add ut + std::shared_ptr pool_; + BTreeFileMetaSelector file_selector_; + std::shared_ptr key_type_; + std::shared_ptr file_reader_; + std::shared_ptr cache_manager_; + std::map> reader_cache_; +}; + +} // namespace paimon diff --git a/src/paimon/common/global_index/global_index_utils.h b/src/paimon/common/global_index/global_index_utils.h new file mode 100644 index 000000000..6f4c5ddc7 --- /dev/null +++ b/src/paimon/common/global_index/global_index_utils.h @@ -0,0 +1,53 @@ +/* + * 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 "arrow/c/abi.h" +#include "fmt/format.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/status.h" +namespace paimon { +class GlobalIndexUtils { + public: + GlobalIndexUtils() = delete; + ~GlobalIndexUtils() = delete; + + static Status CheckRelativeRowIds(::ArrowArray* c_arrow_array, + const std::vector& relative_row_ids, + std::optional expected_next_row_id) { + if (!c_arrow_array) { + return Status::Invalid("CheckRelativeRowIds failed: null c_arrow_array"); + } + int64_t length = c_arrow_array->length; + if (length == 0) { + return Status::OK(); + } + ScopeGuard guard([c_arrow_array]() -> void { ArrowArrayRelease(c_arrow_array); }); + if (static_cast(relative_row_ids.size()) != length) { + return Status::Invalid(fmt::format( + "relative_row_ids length {} mismatch arrow_array length {} in CheckRelativeRowIds", + relative_row_ids.size(), length)); + } + if (expected_next_row_id && relative_row_ids[0] != expected_next_row_id.value()) { + return Status::Invalid(fmt::format( + "first relative_row_ids {} mismatch inner count {} in CheckRelativeRowIds", + relative_row_ids[0], expected_next_row_id.value())); + } + guard.Release(); + return Status::OK(); + } +}; +} // namespace paimon diff --git a/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index.cpp b/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index.cpp index 5d1020b76..8773e412f 100644 --- a/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index.cpp +++ b/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index.cpp @@ -42,9 +42,9 @@ Result> RangeBitmapGlobalIndex::CreateReader( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr reader, index_->CreateReader(arrow_schema, /*start=*/0, meta.file_size, in, pool)); - auto transform = [range_end = meta.range_end](const std::shared_ptr& result) + auto transform = [](const std::shared_ptr& result) -> Result> { - return FileIndexReaderWrapper::ToGlobalIndexResult(range_end, result); + return FileIndexReaderWrapper::ToGlobalIndexResult(result); }; return std::make_shared(reader, transform); } diff --git a/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index_test.cpp b/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index_test.cpp index ee5c915c1..1547ee95f 100644 --- a/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index_test.cpp +++ b/src/paimon/common/global_index/rangebitmap/range_bitmap_global_index_test.cpp @@ -73,14 +73,15 @@ class RangeBitmapGlobalIndexTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array)); + std::vector row_ids(array->length(), 0); + std::iota(row_ids.begin(), row_ids.end(), 0); + PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array, std::move(row_ids))); PAIMON_ASSIGN_OR_RAISE(auto result_metas, global_writer->Finish()); // check meta EXPECT_EQ(result_metas.size(), 1); auto file_name = PathUtil::GetName(result_metas[0].file_path); EXPECT_TRUE(StringUtils::StartsWith(file_name, "range-bitmap-global-index-")); EXPECT_TRUE(StringUtils::EndsWith(file_name, ".index")); - EXPECT_EQ(result_metas[0].range_end, expected_range.to); EXPECT_FALSE(result_metas[0].metadata); return result_metas[0]; } diff --git a/src/paimon/common/global_index/wrap/file_index_reader_wrapper.h b/src/paimon/common/global_index/wrap/file_index_reader_wrapper.h index 91427172c..e87a1fc40 100644 --- a/src/paimon/common/global_index/wrap/file_index_reader_wrapper.h +++ b/src/paimon/common/global_index/wrap/file_index_reader_wrapper.h @@ -138,14 +138,9 @@ class FileIndexReaderWrapper : public GlobalIndexReader { /// Converts a `FileIndexResult` to a `GlobalIndexResult` by mapping 32-bit row IDs /// to 64-bit global row IDs. static Result> ToGlobalIndexResult( - int64_t range_end, const std::shared_ptr& result) { + const std::shared_ptr& result) { if (auto remain = std::dynamic_pointer_cast(result)) { - return std::make_shared( - [range_end]() -> Result { - RoaringBitmap64 bitmap; - bitmap.AddRange(0, range_end + 1); - return bitmap; - }); + return std::shared_ptr(); } else if (auto skip = std::dynamic_pointer_cast(result)) { return std::make_shared( []() -> Result { return RoaringBitmap64(); }); diff --git a/src/paimon/common/global_index/wrap/file_index_reader_wrapper_test.cpp b/src/paimon/common/global_index/wrap/file_index_reader_wrapper_test.cpp index 174fbf834..785b295b1 100644 --- a/src/paimon/common/global_index/wrap/file_index_reader_wrapper_test.cpp +++ b/src/paimon/common/global_index/wrap/file_index_reader_wrapper_test.cpp @@ -40,12 +40,12 @@ TEST(FileIndexReaderWrapperTest, TestToGlobalIndexResult) { { ASSERT_OK_AND_ASSIGN(auto global_result, FileIndexReaderWrapper::ToGlobalIndexResult( - /*range_end=*/5l, FileIndexResult::Remain())); - check_result(global_result, {0l, 1l, 2l, 3l, 4l, 5l}); + FileIndexResult::Remain())); + ASSERT_FALSE(global_result); } { - ASSERT_OK_AND_ASSIGN(auto global_result, FileIndexReaderWrapper::ToGlobalIndexResult( - /*range_end=*/5l, FileIndexResult::Skip())); + ASSERT_OK_AND_ASSIGN(auto global_result, + FileIndexReaderWrapper::ToGlobalIndexResult(FileIndexResult::Skip())); check_result(global_result, {}); } { @@ -53,8 +53,8 @@ TEST(FileIndexReaderWrapperTest, TestToGlobalIndexResult) { return RoaringBitmap32::From({1, 4, 2147483647}); }; auto file_result = std::make_shared(bitmap_supplier); - ASSERT_OK_AND_ASSIGN(auto global_result, FileIndexReaderWrapper::ToGlobalIndexResult( - /*range_end=*/2147483647l, file_result)); + ASSERT_OK_AND_ASSIGN(auto global_result, + FileIndexReaderWrapper::ToGlobalIndexResult(file_result)); check_result(global_result, {1l, 4l, 2147483647l}); } { @@ -68,7 +68,7 @@ TEST(FileIndexReaderWrapperTest, TestToGlobalIndexResult) { }; auto file_result = std::make_shared(); ASSERT_NOK_WITH_MSG( - FileIndexReaderWrapper::ToGlobalIndexResult(/*range_end=*/10l, file_result), + FileIndexReaderWrapper::ToGlobalIndexResult(file_result), "invalid FileIndexResult, supposed to be Remain or Skip or BitmapIndexResult"); } } diff --git a/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h b/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h index f62084e73..e27f6d0e6 100644 --- a/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h +++ b/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h @@ -25,7 +25,9 @@ #include #include "arrow/c/abi.h" +#include "arrow/c/helpers.h" #include "fmt/format.h" +#include "paimon/common/global_index/global_index_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/file_index/file_index_writer.h" #include "paimon/global_index/global_index_writer.h" @@ -40,10 +42,11 @@ class FileIndexWriterWrapper : public GlobalIndexWriter { const std::shared_ptr& writer) : index_type_(index_type), file_manager_(file_manager), writer_(writer) {} - Status AddBatch(::ArrowArray* c_arrow_array) override { - int64_t length = c_arrow_array->length; + Status AddBatch(::ArrowArray* c_arrow_array, std::vector&& relative_row_ids) override { + PAIMON_RETURN_NOT_OK( + GlobalIndexUtils::CheckRelativeRowIds(c_arrow_array, relative_row_ids, count_)); PAIMON_RETURN_NOT_OK(writer_->AddBatch(c_arrow_array)); - count_ += length; + count_ += c_arrow_array->length; return Status::OK(); } @@ -73,7 +76,6 @@ class FileIndexWriterWrapper : public GlobalIndexWriter { PAIMON_RETURN_NOT_OK(out->Flush()); PAIMON_RETURN_NOT_OK(out->Close()); GlobalIndexIOMeta meta(file_manager_->ToPath(file_name), /*file_size=*/bytes->size(), - /*range_end=*/count_ - 1, /*metadata=*/nullptr); return std::vector({meta}); } diff --git a/src/paimon/common/lookup/sort/sort_lookup_store_factory.cpp b/src/paimon/common/lookup/sort/sort_lookup_store_factory.cpp index e81acdf1d..956b45c8a 100644 --- a/src/paimon/common/lookup/sort/sort_lookup_store_factory.cpp +++ b/src/paimon/common/lookup/sort/sort_lookup_store_factory.cpp @@ -36,9 +36,10 @@ Result> SortLookupStoreFactory::CreateReader( const std::shared_ptr& fs, const std::string& file_path, const std::shared_ptr& pool) const { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, fs->Open(file_path)); + auto block_cache = std::make_shared(file_path, in, cache_manager_, pool); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr reader, - SstFileReader::CreateForSortLookupStore(in, comparator_, cache_manager_, pool)); + SstFileReader::CreateForSortLookupStore(in, comparator_, block_cache, pool)); return std::make_unique(in, reader); } diff --git a/src/paimon/common/sst/sst_file_io_test.cpp b/src/paimon/common/sst/sst_file_io_test.cpp index dc1305e9c..a9a3c6023 100644 --- a/src/paimon/common/sst/sst_file_io_test.cpp +++ b/src/paimon/common/sst/sst_file_io_test.cpp @@ -144,8 +144,9 @@ TEST_P(SstFileIOTest, TestSimple) { // test read ASSERT_OK_AND_ASSIGN(in, fs_->Open(index_path)); - ASSERT_OK_AND_ASSIGN(auto reader, SstFileReader::CreateForSortLookupStore( - in, comparator_, cache_manager_, pool_)); + auto block_cache = std::make_shared(index_path, in, cache_manager_, pool_); + ASSERT_OK_AND_ASSIGN( + auto reader, SstFileReader::CreateForSortLookupStore(in, comparator_, block_cache, pool_)); // not exist key std::string k0 = "k0"; @@ -178,8 +179,9 @@ TEST_P(SstFileIOTest, TestJavaCompatibility) { ASSERT_OK_AND_ASSIGN(std::shared_ptr in, fs_->Open(file)); // test read - ASSERT_OK_AND_ASSIGN(auto reader, SstFileReader::CreateForSortLookupStore( - in, comparator_, cache_manager_, pool_)); + auto block_cache = std::make_shared(file, in, cache_manager_, pool_); + ASSERT_OK_AND_ASSIGN( + auto reader, SstFileReader::CreateForSortLookupStore(in, comparator_, block_cache, pool_)); // not exist key std::string k0 = "10000"; ASSERT_FALSE(reader->Lookup(std::make_shared(k0, pool_.get())).value()); @@ -270,8 +272,9 @@ TEST_F(SstFileIOTest, TestIOException) { CHECK_HOOK_STATUS(in_result.status(), i); std::shared_ptr in = std::move(in_result).value(); + auto block_cache = std::make_shared(index_path, in, cache_manager_, pool_); auto reader_result = - SstFileReader::CreateForSortLookupStore(in, comparator_, cache_manager_, pool_); + SstFileReader::CreateForSortLookupStore(in, comparator_, block_cache, pool_); CHECK_HOOK_STATUS(reader_result.status(), i); std::shared_ptr reader = std::move(reader_result).value(); diff --git a/src/paimon/common/sst/sst_file_reader.cpp b/src/paimon/common/sst/sst_file_reader.cpp index 9b2434f4b..b86b5094d 100644 --- a/src/paimon/common/sst/sst_file_reader.cpp +++ b/src/paimon/common/sst/sst_file_reader.cpp @@ -25,11 +25,9 @@ namespace paimon { Result> SstFileReader::Create( const std::shared_ptr& in, const BlockHandle& index_block_handle, const std::optional& bloom_filter_handle, - MemorySlice::SliceComparator comparator, const std::shared_ptr& cache_manager, + MemorySlice::SliceComparator comparator, const std::shared_ptr& block_cache, const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(std::string file_path, in->GetUri()); - auto block_cache = std::make_shared(file_path, in, cache_manager, pool); - // read bloom filter directly now std::shared_ptr bloom_filter = nullptr; if (bloom_filter_handle.has_value() && @@ -67,7 +65,7 @@ Result> SstFileReader::Create( Result> SstFileReader::CreateForSortLookupStore( const std::shared_ptr& in, MemorySlice::SliceComparator comparator, - const std::shared_ptr& cache_manager, const std::shared_ptr& pool) { + const std::shared_ptr& block_cache, const std::shared_ptr& pool) { PAIMON_ASSIGN_OR_RAISE(uint64_t file_len, in->Length()); PAIMON_RETURN_NOT_OK( in->Seek(file_len - SortLookupStoreFooter::ENCODED_LENGTH, SeekOrigin::FS_SEEK_SET)); @@ -80,7 +78,7 @@ Result> SstFileReader::CreateForSortLookupStore( SortLookupStoreFooter::ReadSortLookupStoreFooter(&footer_input)); return SstFileReader::Create(in, read_footer->GetIndexBlockHandle(), read_footer->GetBloomFilterHandle(), std::move(comparator), - cache_manager, pool); + block_cache, pool); } SstFileReader::SstFileReader(const std::shared_ptr& pool, diff --git a/src/paimon/common/sst/sst_file_reader.h b/src/paimon/common/sst/sst_file_reader.h index 982e9e30f..b8ea63555 100644 --- a/src/paimon/common/sst/sst_file_reader.h +++ b/src/paimon/common/sst/sst_file_reader.h @@ -42,7 +42,7 @@ class PAIMON_EXPORT SstFileReader { static Result> Create( const std::shared_ptr& input, const BlockHandle& index_block_handle, const std::optional& bloom_filter_handle, - MemorySlice::SliceComparator comparator, const std::shared_ptr& cache_manager, + MemorySlice::SliceComparator comparator, const std::shared_ptr& block_cache, const std::shared_ptr& pool); /// Create an SstFileReader by reading the SortLookupStoreFooter from the given InputStream. @@ -50,8 +50,11 @@ class PAIMON_EXPORT SstFileReader { /// creating the reader, which avoids code duplication across callers. static Result> CreateForSortLookupStore( const std::shared_ptr& input, MemorySlice::SliceComparator comparator, - const std::shared_ptr& cache_manager, - const std::shared_ptr& pool); + const std::shared_ptr& block_cache, const std::shared_ptr& pool); + + ~SstFileReader() { + [[maybe_unused]] Status _ = Close(); + } /// Create an iterator for the index block. std::unique_ptr CreateIndexIterator(); diff --git a/src/paimon/core/global_index/global_index_evaluator.h b/src/paimon/core/global_index/global_index_evaluator.h index ed2177fe7..05342d2f6 100644 --- a/src/paimon/core/global_index/global_index_evaluator.h +++ b/src/paimon/core/global_index/global_index_evaluator.h @@ -38,11 +38,11 @@ class PAIMON_EXPORT GlobalIndexEvaluator { /// vector similarity scoring is effectively limited to rows that satisfy /// the predicate. /// @return A `Result` containing: - /// - `std::nullopt` if the predicate cannot be evaluated by this index (e.g., field has + /// - `nullptr` if the predicate cannot be evaluated by this index (e.g., field has /// no index), /// - A `std::shared_ptr` if evaluation succeeds. /// The `GlobalIndexResult` indicates the matching rows (e.g., via row ID bitmaps). - virtual Result>> Evaluate( + virtual Result> Evaluate( const std::shared_ptr& predicate, const std::shared_ptr& vector_search) = 0; }; diff --git a/src/paimon/core/global_index/global_index_evaluator_impl.cpp b/src/paimon/core/global_index/global_index_evaluator_impl.cpp index 79b7cff44..12093fba1 100644 --- a/src/paimon/core/global_index/global_index_evaluator_impl.cpp +++ b/src/paimon/core/global_index/global_index_evaluator_impl.cpp @@ -22,10 +22,10 @@ #include "paimon/predicate/predicate_utils.h" namespace paimon { -Result>> GlobalIndexEvaluatorImpl::Evaluate( +Result> GlobalIndexEvaluatorImpl::Evaluate( const std::shared_ptr& predicate, const std::shared_ptr& vector_search) { - std::optional> compound_result; + std::shared_ptr compound_result; if (predicate) { PAIMON_ASSIGN_OR_RAISE(compound_result, EvaluatePredicate(predicate)); } @@ -53,10 +53,9 @@ Result>> GlobalIndexEvaluatorImpl return readers; } -Result>> -GlobalIndexEvaluatorImpl::EvaluateVectorSearch( +Result> GlobalIndexEvaluatorImpl::EvaluateVectorSearch( const std::shared_ptr& vector_search, - const std::optional>& predicate_result) { + const std::shared_ptr& predicate_result) { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, GetIndexReaders(vector_search->field_name)); if (readers.empty()) { @@ -72,7 +71,7 @@ GlobalIndexEvaluatorImpl::EvaluateVectorSearch( auto final_vector_search = vector_search; if (predicate_result) { auto bitmap_global_index_result = - std::dynamic_pointer_cast(predicate_result.value()); + std::dynamic_pointer_cast(predicate_result); if (!bitmap_global_index_result) { return Status::Invalid( "The pre_filter of vector search only supports BitmapGlobalIndexResult"); @@ -87,13 +86,13 @@ GlobalIndexEvaluatorImpl::EvaluateVectorSearch( } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr scored_result, vector_search_reader->VisitVectorSearch(final_vector_search)); - return std::optional>(scored_result); + return scored_result; } -Result>> -GlobalIndexEvaluatorImpl::EvaluatePredicate(const std::shared_ptr& predicate) { +Result> GlobalIndexEvaluatorImpl::EvaluatePredicate( + const std::shared_ptr& predicate) { if (predicate == nullptr) { - return std::optional>(); + return std::shared_ptr(nullptr); } if (auto compound_predicate = std::dynamic_pointer_cast(predicate)) { @@ -103,23 +102,26 @@ GlobalIndexEvaluatorImpl::EvaluatePredicate(const std::shared_ptr& pr PAIMON_ASSIGN_OR_RAISE(std::vector> readers, GetIndexReaders(field_name)); // calculate compound result as field may has multiple indexes - std::optional> compound_result; + std::shared_ptr compound_result; for (const auto& index_reader : readers) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr sub_result, PredicateUtils::VisitPredicate>(leaf_predicate, index_reader)); - if (!compound_result) { - compound_result = sub_result; - } else { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr and_result, - compound_result.value()->And(sub_result)); - compound_result = and_result; + if (sub_result) { + if (!compound_result) { + compound_result = sub_result; + } else { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr and_result, + compound_result->And(sub_result)); + compound_result = and_result; + } } - assert(compound_result); - PAIMON_ASSIGN_OR_RAISE(bool is_empty, compound_result.value()->IsEmpty()); - if (is_empty) { - return compound_result; + if (compound_result) { + PAIMON_ASSIGN_OR_RAISE(bool is_empty, compound_result->IsEmpty()); + if (is_empty) { + return compound_result; + } } } return compound_result; @@ -128,43 +130,42 @@ GlobalIndexEvaluatorImpl::EvaluatePredicate(const std::shared_ptr& pr "cannot cast predicate {} to CompoundPredicate or LeafPredicate", predicate->ToString())); } -Result>> -GlobalIndexEvaluatorImpl::EvaluateCompoundPredicate( +Result> GlobalIndexEvaluatorImpl::EvaluateCompoundPredicate( const std::shared_ptr& compound_predicate) { if (compound_predicate->GetFunction().GetType() == Function::Type::OR) { - std::optional> compound_result; + std::shared_ptr compound_result; for (const auto& child : compound_predicate->Children()) { - PAIMON_ASSIGN_OR_RAISE(std::optional> sub_result, + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr sub_result, EvaluatePredicate(child)); if (!sub_result) { - return std::optional>(); + return std::shared_ptr(nullptr); } if (!compound_result) { compound_result = sub_result; } else { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr or_result, - compound_result.value()->Or(sub_result.value())); + compound_result->Or(sub_result)); compound_result = or_result; } } return compound_result; } else if (compound_predicate->GetFunction().GetType() == Function::Type::AND) { - std::optional> compound_result; + std::shared_ptr compound_result; for (const auto& child : compound_predicate->Children()) { - PAIMON_ASSIGN_OR_RAISE(std::optional> sub_result, + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr sub_result, EvaluatePredicate(child)); if (sub_result) { if (!compound_result) { compound_result = sub_result; } else { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr and_result, - compound_result.value()->And(sub_result.value())); + compound_result->And(sub_result)); compound_result = and_result; } } if (compound_result) { - PAIMON_ASSIGN_OR_RAISE(bool is_empty, compound_result.value()->IsEmpty()); + PAIMON_ASSIGN_OR_RAISE(bool is_empty, compound_result->IsEmpty()); if (is_empty) { return compound_result; } diff --git a/src/paimon/core/global_index/global_index_evaluator_impl.h b/src/paimon/core/global_index/global_index_evaluator_impl.h index 2d3731a16..bf4f1178c 100644 --- a/src/paimon/core/global_index/global_index_evaluator_impl.h +++ b/src/paimon/core/global_index/global_index_evaluator_impl.h @@ -38,19 +38,19 @@ class GlobalIndexEvaluatorImpl : public GlobalIndexEvaluator { IndexReadersCreator create_index_readers) : table_schema_(table_schema), create_index_readers_(std::move(create_index_readers)) {} - Result>> Evaluate( + Result> Evaluate( const std::shared_ptr& predicate, const std::shared_ptr& vector_search) override; private: - Result>> EvaluateVectorSearch( + Result> EvaluateVectorSearch( const std::shared_ptr& vector_search, - const std::optional>& predicate_result); + const std::shared_ptr& predicate_result); - Result>> EvaluatePredicate( + Result> EvaluatePredicate( const std::shared_ptr& predicate); - Result>> EvaluateCompoundPredicate( + Result> EvaluateCompoundPredicate( const std::shared_ptr& compound_predicate); Result>> GetIndexReaders( diff --git a/src/paimon/core/global_index/global_index_scan_impl.cpp b/src/paimon/core/global_index/global_index_scan_impl.cpp index a46c08837..89851b4ae 100644 --- a/src/paimon/core/global_index/global_index_scan_impl.cpp +++ b/src/paimon/core/global_index/global_index_scan_impl.cpp @@ -145,7 +145,7 @@ Status GlobalIndexScanImpl::Scan() { return Status::OK(); } -Result>> GlobalIndexScanImpl::ParallelScan( +Result> GlobalIndexScanImpl::ParallelScan( const std::vector& ranges, const std::shared_ptr& predicate, const std::shared_ptr& vector_search, const std::shared_ptr& executor) { std::vector> range_scanners; @@ -162,23 +162,22 @@ Result>> GlobalIndexScanImpl::P range_scanners.push_back(scanner_impl); } - std::vector>>>> futures; + std::vector>>> futures; for (size_t i = 0; i < range_scanners.size(); i++) { const auto& scanner = range_scanners[i]; const auto& range = ranges[i]; - auto search_index = - [&scanner, &predicate, &vector_search, - &range]() -> Result>> { + auto search_index = [&scanner, &predicate, &vector_search, + &range]() -> Result> { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr evaluator, scanner->CreateIndexEvaluator()); - PAIMON_ASSIGN_OR_RAISE(std::optional> index_result, + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr index_result, evaluator->Evaluate(predicate, vector_search)); if (!index_result) { return index_result; } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result_with_offset, - index_result.value()->AddOffset(range.from)); - return std::optional>(result_with_offset); + index_result->AddOffset(range.from)); + return result_with_offset; }; futures.push_back(Via(executor.get(), search_index)); } @@ -186,32 +185,31 @@ Result>> GlobalIndexScanImpl::P // collect inner result and check all null bool all_null = true; - std::vector>> results; + std::vector> results; for (auto& result : collected_results) { - PAIMON_ASSIGN_OR_RAISE(std::optional> inner_result, - result); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr inner_result, result); if (inner_result) { all_null = false; } results.push_back(std::move(inner_result)); } if (all_null) { - return std::optional>(); + return std::shared_ptr(nullptr); } // union result from multiple ranges - std::optional> final_global_index_result; + std::shared_ptr final_global_index_result; for (size_t i = 0; i < results.size(); ++i) { std::shared_ptr result = - results[i] ? results[i].value() : BitmapGlobalIndexResult::FromRanges({ranges[i]}); + results[i] ? results[i] : BitmapGlobalIndexResult::FromRanges({ranges[i]}); if (!final_global_index_result) { final_global_index_result = result; } else { PAIMON_ASSIGN_OR_RAISE(final_global_index_result, - final_global_index_result.value()->Or(result)); + final_global_index_result->Or(result)); } } - return std::optional>(final_global_index_result); + return final_global_index_result; } } // namespace paimon diff --git a/src/paimon/core/global_index/global_index_scan_impl.h b/src/paimon/core/global_index/global_index_scan_impl.h index 60b44c9b3..1f7e23cd0 100644 --- a/src/paimon/core/global_index/global_index_scan_impl.h +++ b/src/paimon/core/global_index/global_index_scan_impl.h @@ -47,7 +47,7 @@ class GlobalIndexScanImpl : public GlobalIndexScan { return snapshot_; } - Result>> ParallelScan( + Result> ParallelScan( const std::vector& ranges, const std::shared_ptr& predicate, const std::shared_ptr& vector_search, const std::shared_ptr& executor); diff --git a/src/paimon/core/global_index/global_index_write_task.cpp b/src/paimon/core/global_index/global_index_write_task.cpp index c92dbe76b..cd0e3c6e7 100644 --- a/src/paimon/core/global_index/global_index_write_task.cpp +++ b/src/paimon/core/global_index/global_index_write_task.cpp @@ -17,6 +17,7 @@ #include "paimon/global_index/global_index_write_task.h" #include "arrow/c/bridge.h" +#include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/core_options.h" @@ -82,7 +83,7 @@ Result> CreateBatchReader( .WithFileSystem(core_options.GetFileSystem()) .EnablePrefetch(true) .WithMemoryPool(pool) - .SetReadSchema({field_name}); + .SetReadSchema({field_name, SpecialFields::RowId().Name()}); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, @@ -90,7 +91,7 @@ Result> CreateBatchReader( return table_read->CreateReader(indexed_split); } -Result> BuildIndex(const std::string& field_name, +Result> BuildIndex(const std::string& field_name, const Range& range, BatchReader* batch_reader, GlobalIndexWriter* global_index_writer) { while (true) { @@ -103,18 +104,33 @@ Result> BuildIndex(const std::string& field_name, arrow::ImportArray(c_array.get(), c_schema.get())); auto struct_array = std::dynamic_pointer_cast(array); if (!struct_array) { - return Status::Invalid("array read from batch reader is not a struct array"); + return Status::Invalid( + "array read from batch reader is not a struct array in GlobalIndexWriteTask"); + } + auto indexed_array = struct_array->GetFieldByName(field_name); + if (!indexed_array) { + return Status::Invalid(fmt::format( + "read array does not contain {} field in GlobalIndexWriteTask", field_name)); + } + auto row_id_array = struct_array->GetFieldByName(SpecialFields::RowId().Name()); + auto typed_row_id_array = std::dynamic_pointer_cast(row_id_array); + if (!typed_row_id_array) { + return Status::Invalid( + fmt::format("read array does not contain {} field, or it cannot be casted to " + "Int64Array in GlobalIndexWriteTask", + SpecialFields::RowId().Name())); } - arrow::ArrayVector fields = struct_array->fields(); - fields.erase(fields.begin()); - if (fields.empty()) { - return Status::Invalid("array read from batch reader only contains row kind"); + std::vector relative_row_ids; + relative_row_ids.reserve(typed_row_id_array->length()); + for (int64_t i = 0; i < typed_row_id_array->length(); i++) { + relative_row_ids.push_back(typed_row_id_array->Value(i) - range.from); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr new_array, - arrow::StructArray::Make(fields, {field_name})); + arrow::StructArray::Make({indexed_array}, {field_name})); ::ArrowArray c_new_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*new_array, &c_new_array)); - PAIMON_RETURN_NOT_OK(global_index_writer->AddBatch(&c_new_array)); + PAIMON_RETURN_NOT_OK( + global_index_writer->AddBatch(&c_new_array, std::move(relative_row_ids))); } return global_index_writer->Finish(); } @@ -127,20 +143,15 @@ Result> ToCommitMessage( index_file_metas.reserve(global_index_io_metas.size()); bool is_external_path = file_manager->IsExternalPath(); for (const auto& io_meta : global_index_io_metas) { - if (range.Count() != io_meta.range_end + 1) { - return Status::Invalid( - fmt::format("specified range length {} mismatch indexed range length {}", - range.Count(), io_meta.range_end + 1)); - } std::optional external_path; if (is_external_path) { PAIMON_ASSIGN_OR_RAISE(Path path, PathUtil::ToPath(io_meta.file_path)); external_path = path.ToString(); } index_file_metas.push_back(std::make_shared( - index_type, PathUtil::GetName(io_meta.file_path), io_meta.file_size, - io_meta.range_end + 1, /*dv_ranges=*/std::nullopt, external_path, - GlobalIndexMeta(range.from, io_meta.range_end + range.from, field_id, + index_type, PathUtil::GetName(io_meta.file_path), io_meta.file_size, range.Count(), + /*dv_ranges=*/std::nullopt, external_path, + GlobalIndexMeta(range.from, range.to, field_id, /*extra_field_ids=*/std::nullopt, io_meta.metadata))); } DataIncrement data_increment(std::move(index_file_metas)); @@ -200,8 +211,9 @@ Result> GlobalIndexWriteTask::WriteIndex( CreateBatchReader(table_path, field_name, indexed_split, core_options, pool)); // read from data split and write to index writer - PAIMON_ASSIGN_OR_RAISE(std::vector global_index_io_metas, - BuildIndex(field_name, batch_reader.get(), global_index_writer.get())); + PAIMON_ASSIGN_OR_RAISE( + std::vector global_index_io_metas, + BuildIndex(field_name, range, batch_reader.get(), global_index_writer.get())); // generate commit message return ToCommitMessage(index_type, field.Id(), range, global_index_io_metas, diff --git a/src/paimon/core/global_index/row_range_global_index_scanner_impl.cpp b/src/paimon/core/global_index/row_range_global_index_scanner_impl.cpp index b158d26f1..846a3bd53 100644 --- a/src/paimon/core/global_index/row_range_global_index_scanner_impl.cpp +++ b/src/paimon/core/global_index/row_range_global_index_scanner_impl.cpp @@ -132,7 +132,6 @@ GlobalIndexIOMeta RowRangeGlobalIndexScannerImpl::ToGlobalIndexIOMeta( assert(index_file->GetGlobalIndexMeta()); const auto& global_index_meta = index_file->GetGlobalIndexMeta().value(); return {index_file_manager_->ToPath(index_file), index_file->FileSize(), - /*range_end=*/global_index_meta.row_range_end - global_index_meta.row_range_start, global_index_meta.index_meta}; } diff --git a/src/paimon/core/table/source/data_evolution_batch_scan.cpp b/src/paimon/core/table/source/data_evolution_batch_scan.cpp index 41fd2291b..60e6b2084 100644 --- a/src/paimon/core/table/source/data_evolution_batch_scan.cpp +++ b/src/paimon/core/table/source/data_evolution_batch_scan.cpp @@ -41,11 +41,10 @@ Result> DataEvolutionBatchScan::CreatePlan() { std::optional> row_ranges; std::shared_ptr final_global_index_result = global_index_result_; if (!final_global_index_result) { - PAIMON_ASSIGN_OR_RAISE(std::optional> index_result, - EvalGlobalIndex()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr index_result, EvalGlobalIndex()); if (index_result) { - final_global_index_result = index_result.value(); - PAIMON_ASSIGN_OR_RAISE(row_ranges, index_result.value()->ToRanges()); + final_global_index_result = index_result; + PAIMON_ASSIGN_OR_RAISE(row_ranges, index_result->ToRanges()); } } else { PAIMON_ASSIGN_OR_RAISE(row_ranges, final_global_index_result->ToRanges()); @@ -107,14 +106,13 @@ Result> DataEvolutionBatchScan::WrapToIndexedSplits( return std::make_shared(data_plan->SnapshotId(), indexed_splits); } -Result>> DataEvolutionBatchScan::EvalGlobalIndex() - const { +Result> DataEvolutionBatchScan::EvalGlobalIndex() const { auto predicate = batch_scan_->GetNonPartitionPredicate(); if (!predicate && !vector_search_) { - return std::optional>(); + return std::shared_ptr(nullptr); } if (!core_options_.GlobalIndexEnabled()) { - return std::optional>(); + return std::shared_ptr(nullptr); } auto partition_filter = batch_scan_->GetPartitionPredicate(); // TODO(lisizhuo.lsz): support time travel @@ -128,7 +126,7 @@ Result>> DataEvolutionBatchScan } PAIMON_ASSIGN_OR_RAISE(std::vector indexed_row_ranges, index_scan->GetRowRangeList()); if (indexed_row_ranges.empty()) { - return std::optional>(); + return std::shared_ptr(nullptr); } const auto& snapshot = index_scan_impl->GetSnapshot(); const std::optional& next_row_id = snapshot.NextRowId(); @@ -139,20 +137,18 @@ Result>> DataEvolutionBatchScan std::vector non_indexed_row_ranges = Range(0, next_row_id.value() - 1).Exclude(indexed_row_ranges); PAIMON_ASSIGN_OR_RAISE( - std::optional> index_result, + std::shared_ptr index_result, index_scan_impl->ParallelScan(indexed_row_ranges, predicate, vector_search_, executor_)); if (!index_result) { - return std::optional>(); + return std::shared_ptr(nullptr); } - auto index_result_value = std::move(index_result).value(); if (!non_indexed_row_ranges.empty()) { for (const auto& range : non_indexed_row_ranges) { - PAIMON_ASSIGN_OR_RAISE( - index_result_value, - index_result_value->Or(BitmapGlobalIndexResult::FromRanges({range}))); + PAIMON_ASSIGN_OR_RAISE(index_result, + index_result->Or(BitmapGlobalIndexResult::FromRanges({range}))); } } - return std::optional>(index_result_value); + return index_result; } } // namespace paimon diff --git a/src/paimon/core/table/source/data_evolution_batch_scan.h b/src/paimon/core/table/source/data_evolution_batch_scan.h index 450a9e3c4..fad108ea7 100644 --- a/src/paimon/core/table/source/data_evolution_batch_scan.h +++ b/src/paimon/core/table/source/data_evolution_batch_scan.h @@ -42,7 +42,7 @@ class DataEvolutionBatchScan : public AbstractTableScan { Result> WrapToIndexedSplits( const std::shared_ptr& data_plan, const std::vector& row_ranges, const std::map& id_to_score) const; - Result>> EvalGlobalIndex() const; + Result> EvalGlobalIndex() const; private: std::shared_ptr pool_; diff --git a/src/paimon/global_index/lucene/lucene_global_index_reader.cpp b/src/paimon/global_index/lucene/lucene_global_index_reader.cpp index 3f30b89f7..f460f9c58 100644 --- a/src/paimon/global_index/lucene/lucene_global_index_reader.cpp +++ b/src/paimon/global_index/lucene/lucene_global_index_reader.cpp @@ -92,9 +92,8 @@ Result> LuceneGlobalIndexReader::Create write_options, kJiebaTokenizeMode, std::string(kDefaultJiebaTokenizeMode))); } - return std::shared_ptr( - new LuceneGlobalIndexReader(LuceneUtils::StringToWstring(field_name), io_meta.range_end, - searcher, tokenize_mode, jieba)); + return std::shared_ptr(new LuceneGlobalIndexReader( + LuceneUtils::StringToWstring(field_name), searcher, tokenize_mode, jieba)); } catch (const std::exception& e) { return Status::Invalid( fmt::format("create lucene global index reader failed, with {} error.", e.what())); diff --git a/src/paimon/global_index/lucene/lucene_global_index_reader.h b/src/paimon/global_index/lucene/lucene_global_index_reader.h index b69210957..29c28f13a 100644 --- a/src/paimon/global_index/lucene/lucene_global_index_reader.h +++ b/src/paimon/global_index/lucene/lucene_global_index_reader.h @@ -115,12 +115,11 @@ class LuceneGlobalIndexReader : public GlobalIndexReader { } private: - LuceneGlobalIndexReader(const std::wstring& wfield_name, int64_t range_end, + LuceneGlobalIndexReader(const std::wstring& wfield_name, const Lucene::IndexSearcherPtr& searcher, const std::string& tokenize_mode, const std::shared_ptr& jieba) - : range_end_(range_end), - wfield_name_(wfield_name), + : wfield_name_(wfield_name), searcher_(searcher), tokenize_mode_(tokenize_mode), jieba_(jieba) {} @@ -128,7 +127,7 @@ class LuceneGlobalIndexReader : public GlobalIndexReader { std::vector TokenizeQuery(const std::string& query) const; std::shared_ptr CreateAllResult() const { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return nullptr; } Lucene::QueryPtr ConstructMatchQuery( @@ -152,7 +151,6 @@ class LuceneGlobalIndexReader : public GlobalIndexReader { const std::shared_ptr& full_text_search) const noexcept(false); private: - int64_t range_end_; std::wstring wfield_name_; Lucene::IndexSearcherPtr searcher_; std::string tokenize_mode_; diff --git a/src/paimon/global_index/lucene/lucene_global_index_test.cpp b/src/paimon/global_index/lucene/lucene_global_index_test.cpp index 9a7c78d76..74879c1e4 100644 --- a/src/paimon/global_index/lucene/lucene_global_index_test.cpp +++ b/src/paimon/global_index/lucene/lucene_global_index_test.cpp @@ -80,14 +80,15 @@ class LuceneGlobalIndexTest : public ::testing::Test, ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array)); + std::vector row_ids(array->length(), 0); + std::iota(row_ids.begin(), row_ids.end(), 0); + PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array, std::move(row_ids))); PAIMON_ASSIGN_OR_RAISE(auto result_metas, global_writer->Finish()); // check meta EXPECT_EQ(result_metas.size(), 1); auto file_name = PathUtil::GetName(result_metas[0].file_path); EXPECT_TRUE(StringUtils::StartsWith(file_name, "lucene-fts-global-index-")); EXPECT_TRUE(StringUtils::EndsWith(file_name, ".index")); - EXPECT_EQ(result_metas[0].range_end, expected_range.to); EXPECT_TRUE(result_metas[0].metadata); return result_metas[0]; } diff --git a/src/paimon/global_index/lucene/lucene_global_index_writer.cpp b/src/paimon/global_index/lucene/lucene_global_index_writer.cpp index ec8d4a275..5f523b3c6 100644 --- a/src/paimon/global_index/lucene/lucene_global_index_writer.cpp +++ b/src/paimon/global_index/lucene/lucene_global_index_writer.cpp @@ -19,8 +19,10 @@ #include #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "lucene++/FileUtils.h" #include "lucene++/NoLockFactory.h" +#include "paimon/common/global_index/global_index_utils.h" #include "paimon/common/io/data_output_stream.h" #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/path_util.h" @@ -122,7 +124,10 @@ LuceneGlobalIndexWriter::~LuceneGlobalIndexWriter() { } } -Status LuceneGlobalIndexWriter::AddBatch(::ArrowArray* arrow_array) { +Status LuceneGlobalIndexWriter::AddBatch(::ArrowArray* arrow_array, + std::vector&& relative_row_ids) { + PAIMON_RETURN_NOT_OK( + GlobalIndexUtils::CheckRelativeRowIds(arrow_array, relative_row_ids, row_id_)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(arrow_array, arrow_type_)); auto struct_array = std::dynamic_pointer_cast(array); @@ -231,7 +236,6 @@ Result> LuceneGlobalIndexWriter::Finish() { PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(options_, &options_json)); auto meta_bytes = std::make_shared(options_json, pool_.get()); GlobalIndexIOMeta meta(file_writer_->ToPath(index_file_name), file_size, - /*range_end=*/static_cast(row_id_) - 1, /*metadata=*/meta_bytes); return std::vector({meta}); } diff --git a/src/paimon/global_index/lucene/lucene_global_index_writer.h b/src/paimon/global_index/lucene/lucene_global_index_writer.h index 21df08f19..93d54873d 100644 --- a/src/paimon/global_index/lucene/lucene_global_index_writer.h +++ b/src/paimon/global_index/lucene/lucene_global_index_writer.h @@ -48,7 +48,7 @@ class LuceneGlobalIndexWriter : public GlobalIndexWriter { ~LuceneGlobalIndexWriter() override; - Status AddBatch(::ArrowArray* c_arrow_array) override; + Status AddBatch(::ArrowArray* c_arrow_array, std::vector&& relative_row_ids) override; Result> Finish() override; diff --git a/src/paimon/global_index/lumina/lumina_global_index.cpp b/src/paimon/global_index/lumina/lumina_global_index.cpp index 6541e7edf..041ed53e8 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.cpp +++ b/src/paimon/global_index/lumina/lumina_global_index.cpp @@ -19,6 +19,7 @@ #include #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "lumina/api/Dataset.h" #include "lumina/api/LuminaBuilder.h" #include "lumina/api/LuminaSearcher.h" @@ -26,6 +27,7 @@ #include "lumina/core/Constants.h" #include "lumina/core/Status.h" #include "lumina/core/Types.h" +#include "paimon/common/global_index/global_index_utils.h" #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/rapidjson_util.h" #include "paimon/common/utils/string_utils.h" @@ -161,27 +163,17 @@ Result> LuminaGlobalIndex::CreateReader( searcher->Open(std::move(lumina_file_reader), ::lumina::api::IOOptions())); // check meta - PAIMON_RETURN_NOT_OK(CheckLuminaIndexMeta( - searcher->GetMeta(), /*row_count=*/io_meta.range_end + 1, index_info.dimension)); + if (searcher->GetMeta().dim != index_info.dimension) { + return Status::Invalid( + fmt::format("lumina index dimension {} mismatch dimension {} in io meta", + searcher->GetMeta().dim, index_info.dimension)); + } auto searcher_with_filter = std::make_unique<::lumina::extensions::SearchWithFilterExtension>(); PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher->Attach(*searcher_with_filter)); - return std::make_shared(io_meta.range_end, index_info, std::move(searcher), + return std::make_shared(index_info, std::move(searcher), std::move(searcher_with_filter), lumina_pool); } -Status LuminaGlobalIndex::CheckLuminaIndexMeta(const ::lumina::api::LuminaSearcher::IndexInfo& meta, - int64_t row_count, uint32_t dimension) { - if (meta.dim != dimension) { - return Status::Invalid(fmt::format( - "lumina index dimension {} mismatch dimension {} in io meta", meta.dim, dimension)); - } - if (meta.count != static_cast(row_count)) { - return Status::Invalid(fmt::format( - "lumina index row count {} mismatch row count {} in io meta", meta.count, row_count)); - } - return Status::OK(); -} - class LuminaDataset : public ::lumina::api::Dataset { public: LuminaDataset(int64_t element_count, uint32_t dimension, @@ -242,7 +234,10 @@ LuminaIndexWriter::LuminaIndexWriter(const std::string& field_name, io_options_(std::move(io_options)), lumina_options_(lumina_options) {} -Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array) { +Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, + std::vector&& relative_row_ids) { + PAIMON_RETURN_NOT_OK( + GlobalIndexUtils::CheckRelativeRowIds(arrow_array, relative_row_ids, count_)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(arrow_array, arrow_type_)); if (array->null_count() != 0) { @@ -304,18 +299,16 @@ Result> LuminaIndexWriter::Finish() { PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(lumina_options_, &options_json)); auto meta_bytes = std::make_shared(options_json, pool_->GetPaimonPool().get()); GlobalIndexIOMeta meta(file_manager_->ToPath(index_file_name), file_size, - /*range_end=*/count_ - 1, /*metadata=*/meta_bytes); return std::vector({meta}); } LuminaIndexReader::LuminaIndexReader( - int64_t range_end, const LuminaIndexReader::IndexInfo& index_info, + const LuminaIndexReader::IndexInfo& index_info, std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, const std::shared_ptr& pool) - : range_end_(range_end), - index_info_(index_info), + : index_info_(index_info), pool_(pool), searcher_(std::move(searcher)), searcher_with_filter_(std::move(searcher_with_filter)) {} diff --git a/src/paimon/global_index/lumina/lumina_global_index.h b/src/paimon/global_index/lumina/lumina_global_index.h index d2cbc4d0e..ed1a495c2 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.h +++ b/src/paimon/global_index/lumina/lumina_global_index.h @@ -70,10 +70,6 @@ class LuminaGlobalIndex : public GlobalIndexer { const std::vector& files, const std::shared_ptr& pool) const override; - private: - static Status CheckLuminaIndexMeta(const ::lumina::api::LuminaSearcher::IndexInfo& meta, - int64_t row_count, uint32_t dimension); - private: std::map options_; }; @@ -88,7 +84,7 @@ class LuminaIndexWriter : public GlobalIndexWriter { const std::map& lumina_options, const std::shared_ptr& pool); - Status AddBatch(::ArrowArray* arrow_array) override; + Status AddBatch(::ArrowArray* arrow_array, std::vector&& releative_row_ids) override; Result> Finish() override; @@ -114,8 +110,7 @@ class LuminaIndexReader : public GlobalIndexReader { }; LuminaIndexReader( - int64_t range_end, const IndexInfo& index_info, - std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + const IndexInfo& index_info, std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, const std::shared_ptr& pool); @@ -130,66 +125,66 @@ class LuminaIndexReader : public GlobalIndexReader { Result> VisitFullTextSearch( const std::shared_ptr& full_text_search) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitIsNotNull() override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitIsNull() override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitEqual(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitNotEqual(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitLessThan(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitLessOrEqual(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitGreaterThan(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitGreaterOrEqual( const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitIn( const std::vector& literals) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitNotIn( const std::vector& literals) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitStartsWith(const Literal& prefix) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitEndsWith(const Literal& suffix) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitContains(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } Result> VisitLike(const Literal& literal) override { - return BitmapGlobalIndexResult::FromRanges({Range(0, range_end_)}); + return std::shared_ptr(); } bool IsThreadSafe() const override { @@ -203,7 +198,6 @@ class LuminaIndexReader : public GlobalIndexReader { static Result GetIndexInfo(const GlobalIndexIOMeta& io_meta); private: - int64_t range_end_; LuminaIndexReader::IndexInfo index_info_; std::shared_ptr pool_; std::unique_ptr<::lumina::api::LuminaSearcher> searcher_; diff --git a/src/paimon/global_index/lumina/lumina_global_index_test.cpp b/src/paimon/global_index/lumina/lumina_global_index_test.cpp index c2d752286..d87fb707b 100644 --- a/src/paimon/global_index/lumina/lumina_global_index_test.cpp +++ b/src/paimon/global_index/lumina/lumina_global_index_test.cpp @@ -82,14 +82,15 @@ class LuminaGlobalIndexTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array)); + std::vector row_ids(array->length(), 0); + std::iota(row_ids.begin(), row_ids.end(), 0); + PAIMON_RETURN_NOT_OK(global_writer->AddBatch(&c_array, std::move(row_ids))); PAIMON_ASSIGN_OR_RAISE(auto result_metas, global_writer->Finish()); // check meta EXPECT_EQ(result_metas.size(), 1); auto file_name = PathUtil::GetName(result_metas[0].file_path); EXPECT_TRUE(StringUtils::StartsWith(file_name, "lumina-global-index-")); EXPECT_TRUE(StringUtils::EndsWith(file_name, ".index")); - EXPECT_EQ(result_metas[0].range_end, expected_range.to); EXPECT_TRUE(result_metas[0].metadata); return result_metas[0]; } @@ -208,7 +209,7 @@ TEST_F(LuminaGlobalIndexTest, TestSimple) { { // visit equal will return all rows ASSERT_OK_AND_ASSIGN(auto is_null_result, reader->VisitIsNull()); - ASSERT_EQ(is_null_result->ToString(), "{0,1,2,3}"); + ASSERT_FALSE(is_null_result); } } @@ -364,13 +365,6 @@ TEST_F(LuminaGlobalIndexTest, TestInvalidInputs) { CreateGlobalIndexReader(index_root, data_type_, options_, fake_meta), "non-exist-file\' not exists"); } - { - auto fake_meta = meta; - fake_meta.range_end = 50; - ASSERT_NOK_WITH_MSG( - CreateGlobalIndexReader(index_root, data_type_, options_, fake_meta), - "lumina index row count 4 mismatch row count 51 in io meta"); - } { ASSERT_OK_AND_ASSIGN(auto reader, CreateGlobalIndexReader(index_root, data_type_, options_, meta)); diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index bc03ed462..1d9a7b38d 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -344,15 +344,6 @@ TEST_P(GlobalIndexTest, TestWriteIndex) { /*options=*/{}, pool_), "Unknown index type invalid, may not registered"); } - { - // test invalid range mismatch - ASSERT_NOK_WITH_MSG( - GlobalIndexWriteTask::WriteIndex( - table_path, "f0", "bitmap", - std::make_shared(split, std::vector({Range(0, 8)})), - /*options=*/{}, pool_), - "specified range length 9 mismatch indexed range length 8"); - } { // test invalid multiple ranges ASSERT_NOK_WITH_MSG(GlobalIndexWriteTask::WriteIndex( @@ -470,7 +461,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,7}"); + ASSERT_EQ(index_result->ToString(), "{0,7}"); } { // test not equal predicate for f0 @@ -479,7 +470,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{1,2,3,4,5,6}"); + ASSERT_EQ(index_result->ToString(), "{1,2,3,4,5,6}"); } { // test equal predicate for f1 @@ -487,7 +478,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(20)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{4,6,7}"); + ASSERT_EQ(index_result->ToString(), "{4,6,7}"); } { // test equal predicate for f2 @@ -495,7 +486,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(1)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,4,5}"); + ASSERT_EQ(index_result->ToString(), "{0,1,4,5}"); } { // test is null predicate @@ -503,7 +494,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { PredicateBuilder::IsNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{7}"); + ASSERT_EQ(index_result->ToString(), "{7}"); } { // test is not null predicate @@ -511,7 +502,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { PredicateBuilder::IsNotNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,2,3,4,5,6}"); + ASSERT_EQ(index_result->ToString(), "{0,1,2,3,4,5,6}"); } { // test in predicate @@ -521,7 +512,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { Literal(FieldType::STRING, "Lucy", 4)}); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,4,5,7}"); + ASSERT_EQ(index_result->ToString(), "{0,1,4,5,7}"); } { // test not in predicate @@ -531,7 +522,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { Literal(FieldType::STRING, "Lucy", 4)}); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{2,3,6}"); + ASSERT_EQ(index_result->ToString(), "{2,3,6}"); } { // test and predicate @@ -543,7 +534,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_predicate, f1_predicate})); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{7}"); + ASSERT_EQ(index_result->ToString(), "{7}"); } { // test or predicate @@ -555,7 +546,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::Or({f0_predicate, f1_predicate})); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,4,6,7}"); + ASSERT_EQ(index_result->ToString(), "{0,4,6,7}"); } { // test non-result @@ -563,7 +554,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(30)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{}"); + ASSERT_EQ(index_result->ToString(), "{}"); } { // test early stopping @@ -579,7 +570,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { PredicateBuilder::And({f1_predicate, f2_predicate, f0_predicate})); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{}"); + ASSERT_EQ(index_result->ToString(), "{}"); } { // test greater than predicate which bitmap index is not support, will return all range @@ -587,7 +578,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,2,3,4,5,6,7}"); + ASSERT_FALSE(index_result); } { // test greater or equal predicate which bitmap index is not support, will return all range @@ -595,7 +586,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,2,3,4,5,6,7}"); + ASSERT_FALSE(index_result); } { // test less than predicate which bitmap index is not support, will return all range @@ -603,7 +594,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,2,3,4,5,6,7}"); + ASSERT_FALSE(index_result); } { // test less or equal predicate which bitmap index is not support, will return all range @@ -611,7 +602,7 @@ TEST_P(GlobalIndexTest, TestScanIndex) { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,1,2,3,4,5,6,7}"); + ASSERT_FALSE(index_result); } { // test a predicate for field with no index @@ -662,7 +653,7 @@ TEST_P(GlobalIndexTest, TestScanIndexWithSpecificSnapshot) { ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({f0_predicate, f1_predicate})); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,7}"); + ASSERT_EQ(index_result->ToString(), "{0,7}"); } { // test or predicate @@ -749,7 +740,7 @@ TEST_P(GlobalIndexTest, TestScanIndexWithRange) { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN(auto evaluator_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(evaluator_result.value()->ToString(), "{1,2,3,4,5,6}"); + ASSERT_EQ(evaluator_result->ToString(), "{1,2,3,4,5,6}"); } { ASSERT_OK_AND_ASSIGN(auto range_scanner, global_index_scan->CreateRangeScan(Range(10, 13))); @@ -817,7 +808,7 @@ TEST_P(GlobalIndexTest, TestScanIndexWithPartition) { Literal(FieldType::STRING, "Bob", 3)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0,2,3}"); + ASSERT_EQ(index_result->ToString(), "{0,2,3}"); } { // test equal predicate for Alice @@ -826,7 +817,7 @@ TEST_P(GlobalIndexTest, TestScanIndexWithPartition) { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); - ASSERT_EQ(index_result.value()->ToString(), "{0}"); + ASSERT_EQ(index_result->ToString(), "{0}"); } }; @@ -997,7 +988,7 @@ TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndexWithPartition) { ASSERT_OK_AND_ASSIGN(auto index_result, evaluator->Evaluate(predicate, /*vector_search=*/nullptr)); ASSERT_TRUE(index_result); - ASSERT_EQ(index_result.value()->ToString(), bitmap_result); + ASSERT_EQ(index_result->ToString(), bitmap_result); // check lumina index ASSERT_OK_AND_ASSIGN(auto lumina_reader, range_scanner->CreateReader("f1", "lumina")); @@ -1014,13 +1005,13 @@ TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndexWithPartition) { ASSERT_OK_AND_ASSIGN(auto compound_index_result, evaluator->Evaluate(predicate, vector_search_without_filter)); ASSERT_TRUE(compound_index_result); - ASSERT_EQ(compound_index_result.value()->ToString(), lumina_result); + ASSERT_EQ(compound_index_result->ToString(), lumina_result); // check read array std::vector read_field_names = schema->field_names(); read_field_names.push_back("_INDEX_SCORE"); ASSERT_OK_AND_ASSIGN(auto result_with_offset, - compound_index_result.value()->AddOffset(expected_range.from)); + compound_index_result->AddOffset(expected_range.from)); ASSERT_OK_AND_ASSIGN(auto plan, ScanGlobalIndexAndData(table_path, /*predicate=*/nullptr, /*vector_search=*/nullptr, /*options=*/{}, result_with_offset)); diff --git a/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin new file mode 100644 index 000000000..e48af0e1d --- /dev/null +++ b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:df67d2b95ec5d8d6f2f9e3ff1c59965f25ba8b52e13ab9cf0bba21040359f9a7 +size 640 diff --git a/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin.meta b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin.meta new file mode 100644 index 000000000..06f454799 --- /dev/null +++ b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.bin.meta @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:437889433324162e9a2bbc225237b87b33cee8ac4b332c430bf3c44b6eebb96b +size 17 diff --git a/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.csv b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.csv new file mode 100644 index 000000000..f17286d0b --- /dev/null +++ b/test/test_data/global_index/btree/btree_compatibility_data/btree_test_float_50.csv @@ -0,0 +1,51 @@ +row_id,key,is_null +0,-infinity,false +1,NULL,true +2,-9.27f,false +3,-0.00f,false +4,0.00f,false +5,16.41f,false +6,22.79f,false +7,22.79f,false +8,25.32f,false +9,29.65f,false +10,31.18f,false +11,35.94f,false +12,39.07f,false +13,43.51f,false +14,44.23f,false +15,47.69f,false +16,51.82f,false +17,54.36f,false +18,NULL,true +19,61.47f,false +20,62.91f,false +21,64.58f,false +22,69.13f,false +23,73.45f,false +24,75.72f,false +25,78.36f,false +26,78.36f,false +27,85.19f,false +28,86.54f,false +29,91.67f,false +30,92.83f,false +31,93.21f,false +32,100.48f,false +33,102.75f,false +34,104.39f,false +35,106.62f,false +36,108.17f,false +37,115.84f,false +38,115.84f,false +39,117.29f,false +40,121.56f,false +41,126.93f,false +42,128.41f,false +43,131.78f,false +44,135.25f,false +45,138.64f,false +46,+infinity,false +47,NaN,false +48,NULL,true +49,NULL,true From 65e78a3543011f502fcd7cdf8c44b20606ab7c3f Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Mon, 27 Apr 2026 08:26:03 +0000 Subject: [PATCH 2/3] fix comment --- src/paimon/CMakeLists.txt | 2 + .../btree/btree_compatibility_test.cpp | 29 +- .../btree_global_index_integration_test.cpp | 18 - .../btree/btree_global_indexer.cpp | 5 +- .../btree/lazy_filtered_btree_reader.cpp | 144 ++++-- .../btree/lazy_filtered_btree_reader.h | 10 +- .../btree/lazy_filtered_btree_reader_test.cpp | 460 ++++++++++++++++++ .../common/global_index/global_index_utils.h | 18 +- .../global_index/global_index_utils_test.cpp | 108 ++++ .../wrap/file_index_writer_wrapper.h | 3 +- src/paimon/common/sst/sst_file_reader.cpp | 5 +- src/paimon/common/sst/sst_file_reader.h | 2 +- .../global_index/lumina/lumina_global_index.h | 2 +- 13 files changed, 722 insertions(+), 84 deletions(-) create mode 100644 src/paimon/common/global_index/btree/lazy_filtered_btree_reader_test.cpp create mode 100644 src/paimon/common/global_index/global_index_utils_test.cpp diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index d6b3fbc7b..75ac9a727 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -407,6 +407,7 @@ if(PAIMON_BUILD_TESTS) common/file_index/bloomfilter/fast_hash_test.cpp common/global_index/complete_index_score_batch_reader_test.cpp common/global_index/global_index_result_test.cpp + common/global_index/global_index_utils_test.cpp common/global_index/global_indexer_factory_test.cpp common/global_index/bitmap_global_index_result_test.cpp common/global_index/bitmap_scored_global_index_result_test.cpp @@ -417,6 +418,7 @@ if(PAIMON_BUILD_TESTS) common/global_index/btree/btree_global_index_integration_test.cpp common/global_index/btree/btree_compatibility_test.cpp common/global_index/btree/btree_file_meta_selector_test.cpp + common/global_index/btree/lazy_filtered_btree_reader_test.cpp common/global_index/rangebitmap/range_bitmap_global_index_test.cpp common/global_index/wrap/file_index_reader_wrapper_test.cpp common/io/byte_array_input_stream_test.cpp diff --git a/src/paimon/common/global_index/btree/btree_compatibility_test.cpp b/src/paimon/common/global_index/btree/btree_compatibility_test.cpp index 3b58395b8..5079461a1 100644 --- a/src/paimon/common/global_index/btree/btree_compatibility_test.cpp +++ b/src/paimon/common/global_index/btree/btree_compatibility_test.cpp @@ -531,8 +531,7 @@ class BTreeCompatibilityTest : public ::testing::Test { ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotEqual(literal)); auto actual_ids = CollectRowIds(result); auto expected_ids = CollectMatchingRows( - records, - [key_val, rec](const CsvRecord& r) { return !r.is_null && r.key != rec.key; }); + records, [rec](const CsvRecord& r) { return !r.is_null && r.key != rec.key; }); ASSERT_EQ(actual_ids, expected_ids); break; } @@ -718,7 +717,6 @@ TEST_F(BTreeCompatibilityTest, DuplicateKeys) { auto records = ParseCsvFile(csv_path); ASSERT_FALSE(records.empty()); - auto count = static_cast(records.size()); ASSERT_OK_AND_ASSIGN(auto reader, CreateReaderFromFiles(bin_path, meta_path, arrow::int32())); @@ -802,6 +800,31 @@ TEST_F(BTreeCompatibilityTest, MetaDeserialization) { ASSERT_EQ(max_key, Literal(143)); } + // Test float_50 meta + { + std::string meta_path = data_dir_ + "/btree_test_float_50.bin.meta"; + auto meta_str = ReadFileAsString(meta_path); + std::shared_ptr meta_bytes = Bytes::AllocateBytes(meta_str, pool_.get()); + + auto meta = BTreeIndexMeta::Deserialize(meta_bytes, pool_.get()); + ASSERT_TRUE(meta); + + ASSERT_TRUE(meta->HasNulls()); + ASSERT_FALSE(meta->OnlyNulls()); + + ASSERT_TRUE(meta->FirstKey()); + ASSERT_OK_AND_ASSIGN(auto min_key, + KeySerializer::DeserializeKey(MemorySlice::Wrap(meta->FirstKey()), + arrow::float32(), pool_.get())); + ASSERT_EQ(min_key, Literal(static_cast(-INFINITY))); + + ASSERT_TRUE(meta->LastKey()); + ASSERT_OK_AND_ASSIGN(auto max_key, + KeySerializer::DeserializeKey(MemorySlice::Wrap(meta->LastKey()), + arrow::float32(), pool_.get())); + ASSERT_EQ(max_key, Literal(static_cast(std::nan("")))); + } + // Test all_nulls meta { std::string meta_path = data_dir_ + "/btree_test_int_all_nulls.bin.meta"; diff --git a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp index f015988f4..bb94d8467 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp +++ b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp @@ -1654,24 +1654,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, CreateWriterWithNonStructSchema) { "arrow schema must be struct type"); } -TEST_P(BTreeGlobalIndexIntegrationTest, CreateReaderWithMultipleMetas) { - auto file_reader = std::make_shared(fs_, base_path_); - auto field = arrow::field("int_field", arrow::int32()); - auto c_schema = CreateArrowSchema(field); - - std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "128"}, - {BtreeDefs::kBtreeIndexCompression, GetParam()}}; - ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); - - // Provide two fake metas - GlobalIndexIOMeta meta1("fake_path_1", 100, nullptr); - GlobalIndexIOMeta meta2("fake_path_2", 200, nullptr); - std::vector metas = {meta1, meta2}; - - ASSERT_NOK_WITH_MSG(indexer->CreateReader(c_schema.get(), file_reader, metas, pool_), - "exist multiple metas"); -} - TEST_P(BTreeGlobalIndexIntegrationTest, CreateReaderWithMultiFieldSchema) { auto file_reader = std::make_shared(fs_, base_path_); diff --git a/src/paimon/common/global_index/btree/btree_global_indexer.cpp b/src/paimon/common/global_index/btree/btree_global_indexer.cpp index 43c4bfb49..0d5ccd3c2 100644 --- a/src/paimon/common/global_index/btree/btree_global_indexer.cpp +++ b/src/paimon/common/global_index/btree/btree_global_indexer.cpp @@ -33,6 +33,7 @@ #include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/preconditions.h" #include "paimon/core/options/compress_options.h" +#include "paimon/executor.h" #include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/memory/bytes.h" #include "paimon/utils/roaring_bitmap64.h" @@ -96,8 +97,10 @@ Result> BTreeGlobalIndexer::CreateReader( "invalid schema for BTreeGlobalIndexReader, supposed to have single field."); } auto key_type = schema->field(0)->type(); + // TODO(lisizhuo.lsz): Allow users to specify an executor + std::shared_ptr executor = CreateDefaultExecutor(); return std::make_shared(files, key_type, file_reader, cache_manager_, - pool); + pool, executor); } } // namespace paimon diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp index f7dc0efec..6dcb5b148 100644 --- a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp @@ -16,8 +16,10 @@ #include "paimon/common/global_index/btree/lazy_filtered_btree_reader.h" +#include #include +#include "paimon/common/executor/future.h" #include "paimon/common/global_index/btree/btree_file_footer.h" #include "paimon/common/global_index/btree/btree_global_index_reader.h" #include "paimon/common/global_index/btree/btree_index_meta.h" @@ -34,105 +36,121 @@ namespace paimon { LazyFilteredBTreeReader::LazyFilteredBTreeReader( const std::vector& files, const std::shared_ptr& key_type, const std::shared_ptr& file_reader, - const std::shared_ptr& cache_manager, const std::shared_ptr& pool) + const std::shared_ptr& cache_manager, const std::shared_ptr& pool, + const std::shared_ptr& executor) : pool_(pool), file_selector_(files, key_type, pool), key_type_(key_type), file_reader_(file_reader), - cache_manager_(cache_manager) {} + cache_manager_(cache_manager), + executor_(executor) {} Result> LazyFilteredBTreeReader::VisitIsNotNull() { - return DispatchVisit([this]() { return file_selector_.VisitIsNotNull(); }, - [](GlobalIndexReader& reader) { return reader.VisitIsNotNull(); }); + return DispatchVisit( + [this]() { return file_selector_.VisitIsNotNull(); }, + [](const std::shared_ptr& reader) { return reader->VisitIsNotNull(); }); } Result> LazyFilteredBTreeReader::VisitIsNull() { - return DispatchVisit([this]() { return file_selector_.VisitIsNull(); }, - [](GlobalIndexReader& reader) { return reader.VisitIsNull(); }); + return DispatchVisit( + [this]() { return file_selector_.VisitIsNull(); }, + [](const std::shared_ptr& reader) { return reader->VisitIsNull(); }); } Result> LazyFilteredBTreeReader::VisitEqual( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitEqual(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitEqual(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitEqual(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitEqual(literal); + }); } Result> LazyFilteredBTreeReader::VisitNotEqual( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitNotEqual(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitNotEqual(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitNotEqual(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitNotEqual(literal); + }); } Result> LazyFilteredBTreeReader::VisitLessThan( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitLessThan(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitLessThan(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitLessThan(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitLessThan(literal); + }); } Result> LazyFilteredBTreeReader::VisitLessOrEqual( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitLessOrEqual(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitLessOrEqual(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitLessOrEqual(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitLessOrEqual(literal); + }); } Result> LazyFilteredBTreeReader::VisitGreaterThan( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitGreaterThan(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitGreaterThan(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitGreaterThan(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitGreaterThan(literal); + }); } Result> LazyFilteredBTreeReader::VisitGreaterOrEqual( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitGreaterOrEqual(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitGreaterOrEqual(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitGreaterOrEqual(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitGreaterOrEqual(literal); + }); } Result> LazyFilteredBTreeReader::VisitIn( const std::vector& literals) { - return DispatchVisit( - [this, &literals]() { return file_selector_.VisitIn(literals); }, - [&literals](GlobalIndexReader& reader) { return reader.VisitIn(literals); }); + return DispatchVisit([this, &literals]() { return file_selector_.VisitIn(literals); }, + [&literals](const std::shared_ptr& reader) { + return reader->VisitIn(literals); + }); } Result> LazyFilteredBTreeReader::VisitNotIn( const std::vector& literals) { - return DispatchVisit( - [this, &literals]() { return file_selector_.VisitNotIn(literals); }, - [&literals](GlobalIndexReader& reader) { return reader.VisitNotIn(literals); }); + return DispatchVisit([this, &literals]() { return file_selector_.VisitNotIn(literals); }, + [&literals](const std::shared_ptr& reader) { + return reader->VisitNotIn(literals); + }); } Result> LazyFilteredBTreeReader::VisitStartsWith( const Literal& prefix) { - return DispatchVisit( - [this, &prefix]() { return file_selector_.VisitStartsWith(prefix); }, - [&prefix](GlobalIndexReader& reader) { return reader.VisitStartsWith(prefix); }); + return DispatchVisit([this, &prefix]() { return file_selector_.VisitStartsWith(prefix); }, + [&prefix](const std::shared_ptr& reader) { + return reader->VisitStartsWith(prefix); + }); } Result> LazyFilteredBTreeReader::VisitEndsWith( const Literal& suffix) { - return DispatchVisit( - [this, &suffix]() { return file_selector_.VisitEndsWith(suffix); }, - [&suffix](GlobalIndexReader& reader) { return reader.VisitEndsWith(suffix); }); + return DispatchVisit([this, &suffix]() { return file_selector_.VisitEndsWith(suffix); }, + [&suffix](const std::shared_ptr& reader) { + return reader->VisitEndsWith(suffix); + }); } Result> LazyFilteredBTreeReader::VisitContains( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitContains(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitContains(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitContains(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitContains(literal); + }); } Result> LazyFilteredBTreeReader::VisitLike( const Literal& literal) { - return DispatchVisit( - [this, &literal]() { return file_selector_.VisitLike(literal); }, - [&literal](GlobalIndexReader& reader) { return reader.VisitLike(literal); }); + return DispatchVisit([this, &literal]() { return file_selector_.VisitLike(literal); }, + [&literal](const std::shared_ptr& reader) { + return reader->VisitLike(literal); + }); } Result> LazyFilteredBTreeReader::VisitVectorSearch( @@ -147,20 +165,50 @@ Result> LazyFilteredBTreeReader::VisitFullTex Result> LazyFilteredBTreeReader::DispatchVisit( SelectAction select_files, ReaderAction action) { - PAIMON_ASSIGN_OR_RAISE(auto selected_files, select_files()); + PAIMON_ASSIGN_OR_RAISE(std::vector selected_files, select_files()); if (selected_files.empty()) { return std::make_shared([]() { return RoaringBitmap64(); }); } - std::shared_ptr merged_result = nullptr; + // Prepare all readers sequentially (reader_cache_ is not thread-safe) + std::vector> readers; + readers.reserve(selected_files.size()); for (const auto& meta : selected_files) { PAIMON_ASSIGN_OR_RAISE(auto reader, GetOrCreateReader(meta)); - PAIMON_ASSIGN_OR_RAISE(auto result, action(*reader)); + readers.push_back(std::move(reader)); + } + + // Execute actions: parallel if executor is available, sequential otherwise + std::vector>> collected_results; + if (executor_ != nullptr) { + // Parallel: submit all tasks to executor, then collect results in order + std::vector>>> futures; + futures.reserve(readers.size()); + for (const auto& reader : readers) { + futures.push_back( + Via(executor_.get(), + [&action, &reader]() -> Result> { + return action(reader); + })); + } + collected_results = CollectAll(futures); + } else { + // Sequential fallback: execute actions one by one + collected_results.reserve(readers.size()); + for (const auto& reader : readers) { + collected_results.push_back(action(reader)); + } + } + + // Merge results in submission order + std::shared_ptr merged_result = nullptr; + for (auto& result_or_status : collected_results) { + PAIMON_ASSIGN_OR_RAISE(auto result, std::move(result_or_status)); if (result == nullptr) { continue; } if (merged_result == nullptr) { - merged_result = result; + merged_result = std::move(result); } else { PAIMON_ASSIGN_OR_RAISE(merged_result, merged_result->Or(result)); } @@ -225,8 +273,8 @@ Result> LazyFilteredBTreeReader::CreateSingle // Create SST file reader PAIMON_ASSIGN_OR_RAISE( std::shared_ptr sst_file_reader, - SstFileReader::Create(input_stream, footer->GetIndexBlockHandle(), - footer->GetBloomFilterHandle(), comparator, block_cache, pool_)); + SstFileReader::Create(footer->GetIndexBlockHandle(), footer->GetBloomFilterHandle(), + comparator, block_cache, pool_)); return std::make_shared(sst_file_reader, std::move(null_bitmap), min_key, max_key, key_type_, pool_); diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h index bd313a692..372ba7691 100644 --- a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.h @@ -29,6 +29,7 @@ #include "paimon/common/io/cache/cache_manager.h" #include "paimon/common/sst/block_cache.h" #include "paimon/common/sst/block_handle.h" +#include "paimon/executor.h" #include "paimon/global_index/global_index_reader.h" #include "paimon/global_index/io/global_index_file_reader.h" #include "paimon/utils/roaring_bitmap64.h" @@ -40,7 +41,8 @@ class LazyFilteredBTreeReader : public GlobalIndexReader { const std::shared_ptr& key_type, const std::shared_ptr& file_reader, const std::shared_ptr& cache_manager, - const std::shared_ptr& pool); + const std::shared_ptr& pool, + const std::shared_ptr& executor); Result> VisitIsNotNull() override; Result> VisitIsNull() override; @@ -75,8 +77,8 @@ class LazyFilteredBTreeReader : public GlobalIndexReader { private: using SelectAction = std::function>()>; - using ReaderAction = - std::function>(GlobalIndexReader&)>; + using ReaderAction = std::function>( + const std::shared_ptr&)>; Result> DispatchVisit(SelectAction select_files, ReaderAction action); @@ -86,13 +88,13 @@ class LazyFilteredBTreeReader : public GlobalIndexReader { const std::optional& block_handle); private: - // TODO(lisizhuo.lsz): add ut std::shared_ptr pool_; BTreeFileMetaSelector file_selector_; std::shared_ptr key_type_; std::shared_ptr file_reader_; std::shared_ptr cache_manager_; std::map> reader_cache_; + std::shared_ptr executor_; }; } // namespace paimon diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader_test.cpp b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader_test.cpp new file mode 100644 index 000000000..183526ec6 --- /dev/null +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader_test.cpp @@ -0,0 +1,460 @@ +/* + * 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/global_index/btree/lazy_filtered_btree_reader.h" + +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/global_index/btree/btree_global_index_writer.h" +#include "paimon/common/global_index/btree/btree_global_indexer.h" +#include "paimon/executor.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/global_index/io/global_index_file_reader.h" +#include "paimon/global_index/io/global_index_file_writer.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class FakeLazyFileWriter : public GlobalIndexFileWriter { + public: + FakeLazyFileWriter(const std::shared_ptr& fs, const std::string& base_path) + : fs_(fs), base_path_(base_path) {} + + Result NewFileName(const std::string& prefix) const override { + return prefix + "_" + std::to_string(file_counter_++); + } + + Result> NewOutputStream( + const std::string& file_name) const override { + return fs_->Create(base_path_ + "/" + file_name, true); + } + + Result GetFileSize(const std::string& file_name) const override { + PAIMON_ASSIGN_OR_RAISE(auto file_status, fs_->GetFileStatus(base_path_ + "/" + file_name)); + return static_cast(file_status->GetLen()); + } + + std::string ToPath(const std::string& file_name) const override { + return base_path_ + "/" + file_name; + } + + private: + std::shared_ptr fs_; + std::string base_path_; + mutable int64_t file_counter_ = 0; +}; + +class FakeLazyFileReader : public GlobalIndexFileReader { + public: + FakeLazyFileReader(const std::shared_ptr& fs, const std::string& base_path) + : fs_(fs), base_path_(base_path) {} + + Result> GetInputStream( + const std::string& file_path) const override { + return fs_->Open(file_path); + } + + private: + std::shared_ptr fs_; + std::string base_path_; +}; + +class LazyFilteredBTreeReaderTest : public ::testing::Test { + public: + void SetUp() override { + test_dir_ = UniqueTestDirectory::Create("local"); + pool_ = GetDefaultPool(); + fs_ = test_dir_->GetFileSystem(); + base_path_ = test_dir_->Str(); + file_writer_ = std::make_shared(fs_, base_path_); + + // Create 3 btree index files with different key ranges: + // + // file1: keys [1, 1, null, 2, 2], row_ids [0, 1, 2, 3, 4] + // range [1, 2], has_null=true + // + // file2: keys [5, 6, 7], row_ids [5, 6, 7] + // range [5, 7], has_null=false + // + // file3: keys [10, null, 12], row_ids [8, 9, 10] + // range [10, 12], has_null=true + all_metas_ = {}; + WriteSingleFile(R"([[1],[1],[null],[2],[2]])", {0, 1, 2, 3, 4}); + WriteSingleFile(R"([[5],[6],[7]])", {5, 6, 7}); + WriteSingleFile(R"([[10],[null],[12]])", {8, 9, 10}); + } + + void WriteSingleFile(const std::string& json_data, const std::vector& row_ids) { + auto c_schema = CreateArrowSchema(); + std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "4096"}, + {BtreeDefs::kBtreeIndexCompression, "NONE"}}; + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); + ASSERT_OK_AND_ASSIGN( + auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer_, pool_)); + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("int_field", arrow::int32())}), json_data) + .ValueOrDie(); + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + + ASSERT_OK(writer->AddBatch(&c_array, std::vector(row_ids))); + ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); + ASSERT_EQ(metas.size(), 1); + all_metas_.push_back(metas[0]); + } + + std::unique_ptr CreateArrowSchema() const { + auto schema = arrow::schema({arrow::field("int_field", arrow::int32())}); + auto c_schema = std::make_unique(); + EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok()); + return c_schema; + } + + std::shared_ptr CreateReader( + const std::shared_ptr& executor = nullptr) const { + auto file_reader = std::make_shared(fs_, base_path_); + auto cache_manager = std::make_shared(1024 * 1024, 0.5); + return std::make_shared(all_metas_, arrow::int32(), file_reader, + cache_manager, pool_, executor); + } + + void CheckResult(const std::shared_ptr& result, + const std::vector& expected) const { + ASSERT_TRUE(result); + auto typed_result = std::dynamic_pointer_cast(result); + ASSERT_TRUE(typed_result); + ASSERT_OK_AND_ASSIGN(const RoaringBitmap64* bitmap, typed_result->GetBitmap()); + ASSERT_TRUE(bitmap); + ASSERT_EQ(*bitmap, RoaringBitmap64::From(expected)) + << "result=" << bitmap->ToString() + << ", expected=" << RoaringBitmap64::From(expected).ToString(); + } + + void CheckEmpty(const std::shared_ptr& result) const { + ASSERT_TRUE(result); + auto typed_result = std::dynamic_pointer_cast(result); + ASSERT_TRUE(typed_result); + ASSERT_OK_AND_ASSIGN(const RoaringBitmap64* bitmap, typed_result->GetBitmap()); + ASSERT_TRUE(bitmap); + ASSERT_TRUE(bitmap->IsEmpty()); + } + + private: + std::unique_ptr test_dir_; + std::shared_ptr pool_; + std::shared_ptr fs_; + std::string base_path_; + std::shared_ptr file_writer_; + std::vector all_metas_; +}; + +// --- VisitEqual --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitEqualHitSingleFile) { + auto reader = CreateReader(); + // key=1 is in file1 only -> rows 0, 1 + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_1)); + CheckResult(result, {0, 1}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitEqualHitDifferentFile) { + auto reader = CreateReader(); + // key=6 is in file2 only -> row 6 + Literal literal_6(6); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_6)); + CheckResult(result, {6}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitEqualNoMatch) { + auto reader = CreateReader(); + // key=99 is out of all ranges -> empty bitmap + Literal literal_99(99); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_99)); + CheckEmpty(result); +} + +// --- VisitNotEqual --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitNotEqual) { + auto reader = CreateReader(); + // NotEqual 5 -> all non-null rows except row 5 + // non-null rows: 0,1,3,4 (file1) + 5,6,7 (file2) + 8,10 (file3) + // minus row 5 (key=5) -> 0,1,3,4,6,7,8,10 + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotEqual(literal_5)); + CheckResult(result, {0, 1, 3, 4, 6, 7, 8, 10}); +} + +// --- VisitLessThan --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitLessThan) { + auto reader = CreateReader(); + // LessThan 5 -> keys 1,2 from file1 -> rows 0,1,3,4 + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessThan(literal_5)); + CheckResult(result, {0, 1, 3, 4}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitLessThanNoMatch) { + auto reader = CreateReader(); + // LessThan 1 -> no key < 1 -> empty + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessThan(literal_1)); + CheckEmpty(result); +} + +// --- VisitLessOrEqual --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitLessOrEqual) { + auto reader = CreateReader(); + // LessOrEqual 5 -> keys 1,2 from file1 + key 5 from file2 -> rows 0,1,3,4,5 + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessOrEqual(literal_5)); + CheckResult(result, {0, 1, 3, 4, 5}); +} + +// --- VisitGreaterThan --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitGreaterThan) { + auto reader = CreateReader(); + // GreaterThan 7 -> keys 10,12 from file3 -> rows 8,10 + Literal literal_7(7); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterThan(literal_7)); + CheckResult(result, {8, 10}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitGreaterThanNoMatch) { + auto reader = CreateReader(); + // GreaterThan 12 -> no key > 12 -> empty + Literal literal_12(12); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterThan(literal_12)); + CheckEmpty(result); +} + +// --- VisitGreaterOrEqual --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitGreaterOrEqual) { + auto reader = CreateReader(); + // GreaterOrEqual 6 -> keys 6,7 from file2 + keys 10,12 from file3 -> rows 6,7,8,10 + Literal literal_6(6); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterOrEqual(literal_6)); + CheckResult(result, {6, 7, 8, 10}); +} + +// --- VisitIn --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitInAcrossFiles) { + auto reader = CreateReader(); + // IN(2, 7) -> file1 has key 2 (rows 3,4), file2 has key 7 (row 7) + std::vector literals = {Literal(2), Literal(7)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(literals)); + CheckResult(result, {3, 4, 7}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitInNoMatch) { + auto reader = CreateReader(); + // IN(99, 100) -> out of all ranges -> empty + std::vector literals = {Literal(99), Literal(100)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(literals)); + CheckEmpty(result); +} + +// --- VisitNotIn --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitNotIn) { + auto reader = CreateReader(); + // NotIn(1, 5, 10) -> all non-null rows minus rows with key=1,5,10 + // non-null rows: 0,1,3,4 + 5,6,7 + 8,10 + // remove key=1 (rows 0,1), key=5 (row 5), key=10 (row 8) + // -> 3,4,6,7,10 + std::vector literals = {Literal(1), Literal(5), Literal(10)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotIn(literals)); + CheckResult(result, {3, 4, 6, 7, 10}); +} + +// --- VisitIsNull --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitIsNull) { + auto reader = CreateReader(); + // file1 has_null=true (row 2), file3 has_null=true (row 9) + // file2 has no nulls + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNull()); + CheckResult(result, {2, 9}); +} + +// --- VisitIsNotNull --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitIsNotNull) { + auto reader = CreateReader(); + // All non-null rows across all files: 0,1,3,4 + 5,6,7 + 8,10 + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNotNull()); + CheckResult(result, {0, 1, 3, 4, 5, 6, 7, 8, 10}); +} + +// --- VisitVectorSearch / VisitFullTextSearch --- + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitVectorSearchNotSupported) { + auto reader = CreateReader(); + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(nullptr), + "LazyFilteredBTreeReader does not support vector search"); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestVisitFullTextSearchNotSupported) { + auto reader = CreateReader(); + ASSERT_NOK_WITH_MSG(reader->VisitFullTextSearch(nullptr), + "LazyFilteredBTreeReader does not support full text search"); +} + +// --- GetIndexType / IsThreadSafe --- + +TEST_F(LazyFilteredBTreeReaderTest, TestGetIndexType) { + auto reader = CreateReader(); + ASSERT_EQ(reader->GetIndexType(), BtreeDefs::kIdentifier); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestIsNotThreadSafe) { + auto reader = CreateReader(); + ASSERT_FALSE(reader->IsThreadSafe()); +} + +// --- Reader cache verification --- + +TEST_F(LazyFilteredBTreeReaderTest, TestReaderCacheReuse) { + auto reader = CreateReader(); + // Call VisitEqual twice on the same file's range to verify reader cache works + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result1, reader->VisitEqual(literal_1)); + CheckResult(result1, {0, 1}); + + Literal literal_2(2); + ASSERT_OK_AND_ASSIGN(auto result2, reader->VisitEqual(literal_2)); + CheckResult(result2, {3, 4}); +} + +// --- Empty files list --- + +TEST_F(LazyFilteredBTreeReaderTest, TestEmptyFilesList) { + std::vector empty_metas; + auto file_reader = std::make_shared(fs_, base_path_); + auto cache_manager = std::make_shared(1024 * 1024, 0.5); + auto reader = std::make_shared( + empty_metas, arrow::int32(), file_reader, cache_manager, pool_, /*executor=*/nullptr); + + // Any query on empty files should return empty bitmap + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_1)); + CheckEmpty(result); + + ASSERT_OK_AND_ASSIGN(result, reader->VisitIsNotNull()); + CheckEmpty(result); +} + +// ============================================================================ +// Parallel execution with Executor +// ============================================================================ + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitEqual) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // key=1 is in file1 -> rows 0, 1 + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_1)); + CheckResult(result, {0, 1}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitNotEqual) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // NotEqual 5 -> all non-null rows except row 5 -> 0,1,3,4,6,7,8,10 + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotEqual(literal_5)); + CheckResult(result, {0, 1, 3, 4, 6, 7, 8, 10}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitLessThan) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // LessThan 5 -> keys 1,2 from file1 -> rows 0,1,3,4 + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessThan(literal_5)); + CheckResult(result, {0, 1, 3, 4}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitGreaterOrEqual) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // GreaterOrEqual 6 -> keys 6,7 from file2 + keys 10,12 from file3 -> rows 6,7,8,10 + Literal literal_6(6); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterOrEqual(literal_6)); + CheckResult(result, {6, 7, 8, 10}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitIn) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // IN(2, 7) -> file1 has key 2 (rows 3,4), file2 has key 7 (row 7) + std::vector literals = {Literal(2), Literal(7)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(literals)); + CheckResult(result, {3, 4, 7}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitIsNull) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // file1 has_null=true (row 2), file3 has_null=true (row 9) + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNull()); + CheckResult(result, {2, 9}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitIsNotNull) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // All non-null rows across all files: 0,1,3,4 + 5,6,7 + 8,10 + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNotNull()); + CheckResult(result, {0, 1, 3, 4, 5, 6, 7, 8, 10}); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelVisitEqualNoMatch) { + std::shared_ptr executor = CreateDefaultExecutor(); + auto reader = CreateReader(executor); + // key=99 out of range -> empty bitmap + Literal literal_99(99); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_99)); + CheckEmpty(result); +} + +TEST_F(LazyFilteredBTreeReaderTest, TestParallelEmptyFilesList) { + std::shared_ptr executor = CreateDefaultExecutor(); + std::vector empty_metas; + auto file_reader = std::make_shared(fs_, base_path_); + auto cache_manager = std::make_shared(1024 * 1024, 0.5); + auto reader = std::make_shared( + empty_metas, arrow::int32(), file_reader, cache_manager, pool_, executor); + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_1)); + CheckEmpty(result); +} + +} // namespace paimon::test diff --git a/src/paimon/common/global_index/global_index_utils.h b/src/paimon/common/global_index/global_index_utils.h index 6f4c5ddc7..a45eb6edf 100644 --- a/src/paimon/common/global_index/global_index_utils.h +++ b/src/paimon/common/global_index/global_index_utils.h @@ -16,9 +16,11 @@ #pragma once #include "arrow/c/abi.h" +#include "arrow/c/helpers.h" #include "fmt/format.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/status.h" + namespace paimon { class GlobalIndexUtils { public: @@ -32,19 +34,27 @@ class GlobalIndexUtils { return Status::Invalid("CheckRelativeRowIds failed: null c_arrow_array"); } int64_t length = c_arrow_array->length; + ScopeGuard guard([c_arrow_array]() -> void { ArrowArrayRelease(c_arrow_array); }); if (length == 0) { + if (!relative_row_ids.empty()) { + return Status::Invalid( + fmt::format("relative_row_ids length {} mismatch arrow_array length 0 in " + "CheckRelativeRowIds", + relative_row_ids.size())); + } + guard.Release(); return Status::OK(); } - ScopeGuard guard([c_arrow_array]() -> void { ArrowArrayRelease(c_arrow_array); }); if (static_cast(relative_row_ids.size()) != length) { return Status::Invalid(fmt::format( "relative_row_ids length {} mismatch arrow_array length {} in CheckRelativeRowIds", relative_row_ids.size(), length)); } if (expected_next_row_id && relative_row_ids[0] != expected_next_row_id.value()) { - return Status::Invalid(fmt::format( - "first relative_row_ids {} mismatch inner count {} in CheckRelativeRowIds", - relative_row_ids[0], expected_next_row_id.value())); + return Status::Invalid( + fmt::format("first relative_row_ids {} mismatch inner expected_next_row_id {} in " + "CheckRelativeRowIds", + relative_row_ids[0], expected_next_row_id.value())); } guard.Release(); return Status::OK(); diff --git a/src/paimon/common/global_index/global_index_utils_test.cpp b/src/paimon/common/global_index/global_index_utils_test.cpp new file mode 100644 index 000000000..d6d722bed --- /dev/null +++ b/src/paimon/common/global_index/global_index_utils_test.cpp @@ -0,0 +1,108 @@ +/* + * 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/global_index/global_index_utils.h" + +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class GlobalIndexUtilsTest : public ::testing::Test { + public: + /// Helper to create a valid ArrowArray with the given number of int32 elements. + static ArrowArray CreateInt32Array(const std::vector& values) { + arrow::Int32Builder builder; + EXPECT_TRUE(builder.AppendValues(values).ok()); + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + ArrowArray c_array; + EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + return c_array; + } +}; + +TEST_F(GlobalIndexUtilsTest, TestNullArrowArray) { + std::vector row_ids = {0, 1, 2}; + ASSERT_NOK_WITH_MSG(GlobalIndexUtils::CheckRelativeRowIds(nullptr, row_ids, std::nullopt), + "CheckRelativeRowIds failed: null c_arrow_array"); +} + +TEST_F(GlobalIndexUtilsTest, TestEmptyArrayReturnsOk) { + auto c_array = CreateInt32Array({}); + ASSERT_OK(GlobalIndexUtils::CheckRelativeRowIds(&c_array, {}, std::nullopt)); + // ArrowArray was not released by CheckRelativeRowIds (early return for length==0), + // so we release it manually. + if (!ArrowArrayIsReleased(&c_array)) { + ArrowArrayRelease(&c_array); + } +} + +TEST_F(GlobalIndexUtilsTest, TestEmptyArrayReturnsInvalid) { + auto c_array = CreateInt32Array({}); + std::vector row_ids = {0, 1, 2}; + ASSERT_NOK_WITH_MSG( + GlobalIndexUtils::CheckRelativeRowIds(&c_array, row_ids, std::nullopt), + "relative_row_ids length 3 mismatch arrow_array length 0 in CheckRelativeRowIds"); +} + +TEST_F(GlobalIndexUtilsTest, TestMatchingRowIdsAndNoExpectedNextRowId) { + auto c_array = CreateInt32Array({10, 20, 30}); + std::vector row_ids = {0, 1, 2}; + // expected_next_row_id is nullopt, so the first-element check is skipped. + ASSERT_OK(GlobalIndexUtils::CheckRelativeRowIds(&c_array, row_ids, std::nullopt)); + // On success, guard.Release() is called and ArrowArray is NOT released by + // CheckRelativeRowIds, so we must release it manually. + if (!ArrowArrayIsReleased(&c_array)) { + ArrowArrayRelease(&c_array); + } +} + +TEST_F(GlobalIndexUtilsTest, TestMatchingRowIdsWithCorrectExpectedNextRowId) { + auto c_array = CreateInt32Array({10, 20, 30}); + std::vector row_ids = {5, 6, 7}; + // expected_next_row_id matches row_ids[0] + ASSERT_OK(GlobalIndexUtils::CheckRelativeRowIds(&c_array, row_ids, 5)); + // On success, guard.Release() is called and ArrowArray is NOT released by + // CheckRelativeRowIds, so we must release it manually. + if (!ArrowArrayIsReleased(&c_array)) { + ArrowArrayRelease(&c_array); + } +} + +TEST_F(GlobalIndexUtilsTest, TestMismatchedLength) { + auto c_array = CreateInt32Array({10, 20, 30}); + std::vector row_ids = {0, 1}; + ASSERT_NOK_WITH_MSG( + GlobalIndexUtils::CheckRelativeRowIds(&c_array, row_ids, std::nullopt), + "relative_row_ids length 2 mismatch arrow_array length 3 in CheckRelativeRowIds"); +} + +TEST_F(GlobalIndexUtilsTest, TestMismatchedExpectedNextRowId) { + auto c_array = CreateInt32Array({10, 20, 30}); + std::vector row_ids = {5, 6, 7}; + // expected_next_row_id is 100, but row_ids[0] is 5 + ASSERT_NOK_WITH_MSG( + GlobalIndexUtils::CheckRelativeRowIds(&c_array, row_ids, 100), + "first relative_row_ids 5 mismatch inner expected_next_row_id 100 in CheckRelativeRowIds"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h b/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h index e27f6d0e6..e4cc0b9f3 100644 --- a/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h +++ b/src/paimon/common/global_index/wrap/file_index_writer_wrapper.h @@ -45,8 +45,9 @@ class FileIndexWriterWrapper : public GlobalIndexWriter { Status AddBatch(::ArrowArray* c_arrow_array, std::vector&& relative_row_ids) override { PAIMON_RETURN_NOT_OK( GlobalIndexUtils::CheckRelativeRowIds(c_arrow_array, relative_row_ids, count_)); + auto length = c_arrow_array->length; PAIMON_RETURN_NOT_OK(writer_->AddBatch(c_arrow_array)); - count_ += c_arrow_array->length; + count_ += length; return Status::OK(); } diff --git a/src/paimon/common/sst/sst_file_reader.cpp b/src/paimon/common/sst/sst_file_reader.cpp index b86b5094d..ad430ceb9 100644 --- a/src/paimon/common/sst/sst_file_reader.cpp +++ b/src/paimon/common/sst/sst_file_reader.cpp @@ -23,11 +23,10 @@ namespace paimon { Result> SstFileReader::Create( - const std::shared_ptr& in, const BlockHandle& index_block_handle, + const BlockHandle& index_block_handle, const std::optional& bloom_filter_handle, MemorySlice::SliceComparator comparator, const std::shared_ptr& block_cache, const std::shared_ptr& pool) { - PAIMON_ASSIGN_OR_RAISE(std::string file_path, in->GetUri()); // read bloom filter directly now std::shared_ptr bloom_filter = nullptr; if (bloom_filter_handle.has_value() && @@ -76,7 +75,7 @@ Result> SstFileReader::CreateForSortLookupStore( auto footer_input = footer_slice.ToInput(); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_footer, SortLookupStoreFooter::ReadSortLookupStoreFooter(&footer_input)); - return SstFileReader::Create(in, read_footer->GetIndexBlockHandle(), + return SstFileReader::Create(read_footer->GetIndexBlockHandle(), read_footer->GetBloomFilterHandle(), std::move(comparator), block_cache, pool); } diff --git a/src/paimon/common/sst/sst_file_reader.h b/src/paimon/common/sst/sst_file_reader.h index b8ea63555..e115054c4 100644 --- a/src/paimon/common/sst/sst_file_reader.h +++ b/src/paimon/common/sst/sst_file_reader.h @@ -40,7 +40,7 @@ class SstFileIterator; class PAIMON_EXPORT SstFileReader { public: static Result> Create( - const std::shared_ptr& input, const BlockHandle& index_block_handle, + const BlockHandle& index_block_handle, const std::optional& bloom_filter_handle, MemorySlice::SliceComparator comparator, const std::shared_ptr& block_cache, const std::shared_ptr& pool); diff --git a/src/paimon/global_index/lumina/lumina_global_index.h b/src/paimon/global_index/lumina/lumina_global_index.h index ed1a495c2..0c5a11ab5 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.h +++ b/src/paimon/global_index/lumina/lumina_global_index.h @@ -84,7 +84,7 @@ class LuminaIndexWriter : public GlobalIndexWriter { const std::map& lumina_options, const std::shared_ptr& pool); - Status AddBatch(::ArrowArray* arrow_array, std::vector&& releative_row_ids) override; + Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) override; Result> Finish() override; From 65469b00d741e01d0a9f37f25434791dc29e02b1 Mon Sep 17 00:00:00 2001 From: "lisizhuo.lsz" Date: Mon, 27 Apr 2026 08:58:01 +0000 Subject: [PATCH 3/3] mv auto & add ut in btree_global_index_integration_test.cpp && fix global_index_utils --- .../btree_global_index_integration_test.cpp | 244 ++++++++++++------ .../btree/lazy_filtered_btree_reader.cpp | 7 +- .../common/global_index/global_index_utils.h | 16 +- 3 files changed, 176 insertions(+), 91 deletions(-) diff --git a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp index bb94d8467..6454a0045 100644 --- a/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp +++ b/src/paimon/common/global_index/btree/btree_global_index_integration_test.cpp @@ -127,8 +127,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadIntData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); // Data layout (row_id -> value): // 0->1, 1->1, 2->null, 3->2, 4->2, 5->null, 6->3, 7->4, 8->5, 9->5, 10->5, 11->null auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -151,7 +149,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -281,9 +279,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadStringData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("str_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->"apple", 1->"apricot", 2->null, 3->"banana", 4->"blueberry", // 5->null, 6->"cherry", 7->"cherry", 8->"date" @@ -304,7 +299,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadStringData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -449,9 +444,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBigIntData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("bigint_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->100, 1->null, 2->200, 3->200, 4->300, 5->null, 6->400 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -469,7 +461,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBigIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -541,9 +533,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadFloatData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("float_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->1.0, 1->null, 2->2.5, 3->2.5, 4->3.0, 5->null, 6->4.5 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -561,7 +550,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadFloatData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -632,9 +621,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDoubleData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("double_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->1.1, 1->null, 2->2.2, 3->2.2, 4->3.3, 5->null, 6->4.4 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -652,7 +638,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDoubleData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -723,9 +709,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNonNull) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // All values are non-null // Data layout (row_id -> value): 0->10, 1->20, 2->20, 3->30, 4->40, 5->50 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -742,7 +725,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNonNull) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -832,9 +815,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBoolData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("bool_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value), sorted by key (false < true): // 0->false, 1->false, 2->null, 3->false, 4->true, 5->null, 6->true, 7->true auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -853,7 +833,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadBoolData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -909,9 +889,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTinyIntData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("tinyint_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->-10, 1->null, 2->0, 3->10, 4->10, 5->null, 6->20 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -929,7 +906,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTinyIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -991,9 +968,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadSmallIntData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("smallint_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value): // 0->-100, 1->null, 2->0, 3->100, 4->100, 5->null, 6->200 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -1011,7 +985,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadSmallIntData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1074,9 +1048,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampCompactData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("ts_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value in millis): // 0->1000, 1->null, 2->2000, 3->2000, 4->3000, 5->null, 6->4000 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ @@ -1094,7 +1065,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1157,9 +1128,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampNonCompactData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("ts_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value in micros): // 0->1000000 (1s), 1->null, 2->2000123 (2s+123us), 3->2000123, 4->3000456, 5->null, // 6->4000789 @@ -1178,7 +1146,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadTimestampNonCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1245,9 +1213,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalCompactData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("decimal_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> value, stored as unscaled int with scale=2): // 0->1.00, 1->null, 2->2.50, 3->2.50, 4->3.00, // 5->null, 6->4.50 @@ -1266,7 +1231,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1329,9 +1294,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalNonCompactData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN( auto writer, indexer->CreateWriter("decimal_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Data layout (row_id -> unscaled value with scale=3): // 0->1.000, 1->null, 2->2.500, 3->2.500, // 4->3.000, 5->null, 6->4.500 @@ -1351,7 +1313,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadDecimalNonCompactData) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1413,9 +1375,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNull) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // All values are null auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ [null], @@ -1429,7 +1388,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadAllNull) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids(array->length()); std::iota(row_ids.begin(), row_ids.end(), 0); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(row_ids))); ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); ASSERT_EQ(metas.size(), 1); @@ -1501,9 +1460,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadLargeDataWithSmallBlocks) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Generate 50000 sorted int values with some nulls and duplicates. // Pattern: every 100th row is null, values increase by 1 every 3 rows (duplicates). // Data is written in multiple batches of 1000 rows each. @@ -1544,7 +1500,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadLargeDataWithSmallBlocks) { ArrowArray c_array; ASSERT_TRUE(arrow::ExportArray(*struct_array, &c_array).ok()); - ASSERT_OK(btree_writer->AddBatch(&c_array, std::move(batch_row_ids))); + ASSERT_OK(writer->AddBatch(&c_array, std::move(batch_row_ids))); } ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); @@ -1695,11 +1651,8 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNullArray) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - std::vector row_ids = {0, 1, 2}; - ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(nullptr, std::move(row_ids)), + ASSERT_NOK_WITH_MSG(writer->AddBatch(nullptr, std::move(row_ids)), "CheckRelativeRowIds failed: null c_arrow_array"); } @@ -1713,9 +1666,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithMismatchedRowIds) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ [1], [2], @@ -1729,7 +1679,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithMismatchedRowIds) { // Provide wrong number of row_ids (2 instead of 3) std::vector row_ids = {0, 1}; ASSERT_NOK_WITH_MSG( - btree_writer->AddBatch(&c_array, std::move(row_ids)), + writer->AddBatch(&c_array, std::move(row_ids)), "relative_row_ids length 2 mismatch arrow_array length 3 in CheckRelativeRowIds"); } @@ -1743,9 +1693,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNonMonotonicKeys) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Write decreasing keys: 3, 2, 1 auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ [3], @@ -1757,7 +1704,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, AddBatchWithNonMonotonicKeys) { ArrowArray c_array; ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids = {0, 1, 2}; - ASSERT_NOK_WITH_MSG(btree_writer->AddBatch(&c_array, std::move(row_ids)), + ASSERT_NOK_WITH_MSG(writer->AddBatch(&c_array, std::move(row_ids)), "Users must keep written keys monotonically incremental"); } @@ -1771,9 +1718,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, FinishWithEmptyData) { ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); ASSERT_OK_AND_ASSIGN(auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); - auto btree_writer = std::dynamic_pointer_cast(writer); - ASSERT_TRUE(btree_writer); - // Finish without adding any data ASSERT_NOK_WITH_MSG(writer->Finish(), "Should never write an empty btree index file"); } @@ -1802,7 +1746,6 @@ TEST_P(BTreeGlobalIndexIntegrationTest, TestIOException) { auto writer_result = indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_); CHECK_HOOK_STATUS(writer_result.status(), i); auto writer = std::move(writer_result).value(); - auto btree_writer = std::dynamic_pointer_cast(writer); auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), R"([ [1], [2], [null], [3], [4], [5] @@ -1812,7 +1755,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, TestIOException) { ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); std::vector row_ids = {0, 1, 2, 3, 4, 5}; - CHECK_HOOK_STATUS(btree_writer->AddBatch(&c_array, std::move(row_ids)), i); + CHECK_HOOK_STATUS(writer->AddBatch(&c_array, std::move(row_ids)), i); auto finish_result = writer->Finish(); CHECK_HOOK_STATUS(finish_result.status(), i); auto metas = std::move(finish_result).value(); @@ -1841,6 +1784,155 @@ TEST_P(BTreeGlobalIndexIntegrationTest, TestIOException) { ASSERT_TRUE(run_complete); } +// Multiple files with BTreeFileMetaSelector filtering +TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadMultiFilesWithMetaSelector) { + auto file_writer = std::make_shared(fs_, base_path_); + auto field = arrow::field("int_field", arrow::int32()); + + std::map options = {{BtreeDefs::kBtreeIndexBlockSize, "4096"}, + {BtreeDefs::kBtreeIndexCompression, GetParam()}}; + ASSERT_OK_AND_ASSIGN(auto indexer, BTreeGlobalIndexer::Create(options)); + + // Write 3 separate btree index files with non-overlapping key ranges: + // + // file0: keys [1, 2, null, 5], row_ids [0, 1, 2, 3] + // range [1, 5], has_null=true + // + // file1: keys [10, 12, 15], row_ids [4, 5, 6] + // range [10, 15], has_null=false + // + // file2: keys [20, null, 25], row_ids [7, 8, 9] + // range [20, 25], has_null=true + std::vector all_metas; + + auto write_file = [&](const std::string& json_data, const std::vector& row_ids) { + auto c_schema = CreateArrowSchema(field); + ASSERT_OK_AND_ASSIGN( + auto writer, indexer->CreateWriter("int_field", c_schema.get(), file_writer, pool_)); + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), json_data) + .ValueOrDie(); + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + ASSERT_OK(writer->AddBatch(&c_array, std::vector(row_ids))); + ASSERT_OK_AND_ASSIGN(auto metas, writer->Finish()); + ASSERT_EQ(metas.size(), 1); + all_metas.push_back(metas[0]); + }; + + write_file(R"([[1],[2],[null],[5]])", {0, 1, 2, 3}); + write_file(R"([[10],[12],[15]])", {4, 5, 6}); + write_file(R"([[20],[null],[25]])", {7, 8, 9}); + + ASSERT_EQ(all_metas.size(), 3); + + // Create reader over all 3 files (internally uses LazyFilteredBTreeReader + + // BTreeFileMetaSelector) + auto file_reader = std::make_shared(fs_, base_path_); + auto c_schema = CreateArrowSchema(field); + ASSERT_OK_AND_ASSIGN(auto reader, + indexer->CreateReader(c_schema.get(), file_reader, all_metas, pool_)); + + // --- VisitEqual: key=12 -> only file1 is selected by meta selector -> row 5 + { + Literal literal_12(12); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_12)); + CheckResult(result, {5}); + } + + // --- VisitEqual: key=1 -> only file0 selected -> row 0 + { + Literal literal_1(1); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_1)); + CheckResult(result, {0}); + } + + // --- VisitEqual: key=25 -> only file2 selected -> row 9 + { + Literal literal_25(25); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_25)); + CheckResult(result, {9}); + } + + // --- VisitEqual: key=8 -> between file0 and file1 ranges, no file matches -> empty + { + Literal literal_8(8); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitEqual(literal_8)); + CheckResult(result, {}); + } + + // --- VisitLessThan: key < 10 -> only file0 selected -> non-null rows 0,1,3 + { + Literal literal_10(10); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessThan(literal_10)); + CheckResult(result, {0, 1, 3}); + } + + // --- VisitGreaterThan: key > 15 -> only file2 selected -> non-null rows 7,9 + { + Literal literal_15(15); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterThan(literal_15)); + CheckResult(result, {7, 9}); + } + + // --- VisitGreaterOrEqual: key >= 5 -> file0 (key=5) + file1 + file2 -> rows 3,4,5,6,7,9 + { + Literal literal_5(5); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterOrEqual(literal_5)); + CheckResult(result, {3, 4, 5, 6, 7, 9}); + } + + // --- VisitLessOrEqual: key <= 2 -> only file0 selected -> rows 0,1 + { + Literal literal_2(2); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitLessOrEqual(literal_2)); + CheckResult(result, {0, 1}); + } + + // --- VisitIn: IN(2, 15, 25) -> file0 (key=2) + file1 (key=15) + file2 (key=25) -> rows 1,6,9 + { + std::vector literals = {Literal(2), Literal(15), Literal(25)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(literals)); + CheckResult(result, {1, 6, 9}); + } + + // --- VisitIn: IN(8, 17) -> no file contains these keys -> empty + { + std::vector literals = {Literal(8), Literal(17)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIn(literals)); + CheckResult(result, {}); + } + + // --- VisitNotEqual: key != 10 -> all files selected, all non-null rows except key=10 + // non-null: file0 {0,1,3} + file1 {4,5,6} + file2 {7,9} + // minus key=10 (row 4) -> {0,1,3,5,6,7,9} + { + Literal literal_10(10); + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotEqual(literal_10)); + CheckResult(result, {0, 1, 3, 5, 6, 7, 9}); + } + + // --- VisitIsNull: file0 has null (row 2), file2 has null (row 8) -> {2, 8} + { + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNull()); + CheckResult(result, {2, 8}); + } + + // --- VisitIsNotNull: all non-null rows -> {0,1,3,4,5,6,7,9} + { + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitIsNotNull()); + CheckResult(result, {0, 1, 3, 4, 5, 6, 7, 9}); + } + + // --- VisitNotIn: NOT IN(1, 12, 20) -> all files selected + // non-null rows: {0,1,3,4,5,6,7,9} + // minus key=1 (row 0), key=12 (row 5), key=20 (row 7) -> {1,3,4,6,9} + { + std::vector literals = {Literal(1), Literal(12), Literal(20)}; + ASSERT_OK_AND_ASSIGN(auto result, reader->VisitNotIn(literals)); + CheckResult(result, {1, 3, 4, 6, 9}); + } +} + INSTANTIATE_TEST_SUITE_P(Compression, BTreeGlobalIndexIntegrationTest, ::testing::ValuesIn(std::vector({"none", "zstd", "lz4"}))); diff --git a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp index 6dcb5b148..c5f3bdf1f 100644 --- a/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp +++ b/src/paimon/common/global_index/btree/lazy_filtered_btree_reader.cpp @@ -174,7 +174,7 @@ Result> LazyFilteredBTreeReader::DispatchVisi std::vector> readers; readers.reserve(selected_files.size()); for (const auto& meta : selected_files) { - PAIMON_ASSIGN_OR_RAISE(auto reader, GetOrCreateReader(meta)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, GetOrCreateReader(meta)); readers.push_back(std::move(reader)); } @@ -203,7 +203,8 @@ Result> LazyFilteredBTreeReader::DispatchVisi // Merge results in submission order std::shared_ptr merged_result = nullptr; for (auto& result_or_status : collected_results) { - PAIMON_ASSIGN_OR_RAISE(auto result, std::move(result_or_status)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, + std::move(result_or_status)); if (result == nullptr) { continue; } @@ -226,7 +227,7 @@ Result> LazyFilteredBTreeReader::GetOrCreateR if (iterator != reader_cache_.end()) { return iterator->second; } - PAIMON_ASSIGN_OR_RAISE(auto reader, CreateSingleReader(meta)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, CreateSingleReader(meta)); reader_cache_[meta.file_path] = reader; return reader; } diff --git a/src/paimon/common/global_index/global_index_utils.h b/src/paimon/common/global_index/global_index_utils.h index a45eb6edf..aadefc20f 100644 --- a/src/paimon/common/global_index/global_index_utils.h +++ b/src/paimon/common/global_index/global_index_utils.h @@ -15,12 +15,13 @@ */ #pragma once +#include + #include "arrow/c/abi.h" #include "arrow/c/helpers.h" #include "fmt/format.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/status.h" - namespace paimon { class GlobalIndexUtils { public: @@ -35,22 +36,13 @@ class GlobalIndexUtils { } int64_t length = c_arrow_array->length; ScopeGuard guard([c_arrow_array]() -> void { ArrowArrayRelease(c_arrow_array); }); - if (length == 0) { - if (!relative_row_ids.empty()) { - return Status::Invalid( - fmt::format("relative_row_ids length {} mismatch arrow_array length 0 in " - "CheckRelativeRowIds", - relative_row_ids.size())); - } - guard.Release(); - return Status::OK(); - } if (static_cast(relative_row_ids.size()) != length) { return Status::Invalid(fmt::format( "relative_row_ids length {} mismatch arrow_array length {} in CheckRelativeRowIds", relative_row_ids.size(), length)); } - if (expected_next_row_id && relative_row_ids[0] != expected_next_row_id.value()) { + if (!relative_row_ids.empty() && expected_next_row_id && + relative_row_ids[0] != expected_next_row_id.value()) { return Status::Invalid( fmt::format("first relative_row_ids {} mismatch inner expected_next_row_id {} in " "CheckRelativeRowIds",