diff --git a/docs/source/user_guide.rst b/docs/source/user_guide.rst index f0455343b..38ef0fc64 100644 --- a/docs/source/user_guide.rst +++ b/docs/source/user_guide.rst @@ -42,3 +42,4 @@ User Guide user_guide/prefetch user_guide/arrow user_guide/global_index + user_guide/primary_key_global_index diff --git a/docs/source/user_guide/primary_key_global_index.rst b/docs/source/user_guide/primary_key_global_index.rst new file mode 100644 index 000000000..3b9aaa56e --- /dev/null +++ b/docs/source/user_guide/primary_key_global_index.rst @@ -0,0 +1,73 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +Primary Key Global Index +======================== + +Paimon 2.0 primary-key tables support *source-backed* global scalar indexes +(``pk-btree`` / ``pk-bitmap``). Unlike the Data Evolution global indexes described in +:doc:`global_index`, which address rows by a table-wide row id, a source-backed payload +covers the complete active source set of one positive data level of one bucket, and its +results are group ordinals that are localized back to per-file physical row positions. + +paimon-cpp supports the read path of this protocol: ordinary batch scans of a +primary-key table with scalar index definitions automatically evaluate the part of the +scan predicate that touches indexed fields against the validated payload groups of the +scanned snapshot, and narrow covered files to indexed splits carrying file-local row +ranges. No dedicated query API is required. + +Table requirements +------------------ + +The definitions follow the Java table options: + +- ``'pk-btree.index.columns' = 'price'`` with optional + ``'fields.price.pk-btree.index.options' = '{"block-size":"64 kb"}'`` +- fixed bucket (``bucket > 0``) or postpone bucket mode +- ``'deletion-vectors.enabled' = 'true'`` and ``'deletion-vectors.merge-on-read' = 'false'`` + +Semantics +--------- + +- A payload is only used when it provably covers the current active source set of its + data level: exactly one payload per level, source file names / order / row counts + identical to the active COMPACT files of that level, matching index type and field id, + and a row range of exactly ``[0, total source rows - 1]``. Anything else is treated as + uncovered and scanned normally. +- ``AND`` predicates narrow with any safely evaluable indexed child; ``OR`` predicates + only use the index when every branch is evaluable. Files whose evaluation fails, whose + positions are out of range, or whose result needs more than 4096 ranges fall back to a + normal scan individually. +- Indexed splits keep their deletion files aligned with the data file; the reader still + applies deletion vectors and the complete original predicate, so index results never + change visibility semantics. +- ``'global-index.enabled' = 'false'`` disables the planner. + +Current scope +------------- + +- The BTree payload reader is wired up. ``pk-bitmap`` (and vector / full-text) + definitions are recognized for validation, but their evaluation conservatively falls + back to a normal scan until their dedicated payload readers are supported. +- The read path targets the Java release-2.0.0 layout and scan semantics (source metadata + v1, ``GlobalIndexMeta`` with ``_SOURCE_META``, commit message v12). Source-file names + currently use the existing C++ length-prefixed UTF-8 streams; ASCII and non-null BMP + names are compatible with Java ``writeUTF``, while complete modified UTF-8 support for + supplementary code points will be handled by a shared stream-level change. +- ``PkSortedIndexFile::Build`` can build one payload for an ordered source group from + value-sorted input, which supports tooling and tests; automatic build and maintenance + during compaction is not included yet. diff --git a/include/paimon/defs.h b/include/paimon/defs.h index cf01a1ba4..5667ae8fd 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -522,6 +522,18 @@ struct PAIMON_EXPORT Options { /// "global-index.external-path" - Global index root directory, if not set, the global index /// files will be stored under the index directory. static const char GLOBAL_INDEX_EXTERNAL_PATH[]; + /// "pk-btree.index.columns" - Comma-separated columns indexed by primary-key BTree indexes. + /// No default value. + static const char PK_BTREE_INDEX_COLUMNS[]; + /// "pk-bitmap.index.columns" - Comma-separated columns indexed by primary-key Bitmap indexes. + /// No default value. + static const char PK_BITMAP_INDEX_COLUMNS[]; + /// "pk-vector.index.columns" - Comma-separated VECTOR columns indexed by primary-key vector + /// indexes. No default value. + static const char PK_VECTOR_INDEX_COLUMNS[]; + /// "pk-full-text.index.columns" - Comma-separated character columns indexed by primary-key + /// full-text indexes. No default value. + static const char PK_FULL_TEXT_INDEX_COLUMNS[]; /// "aggregation.remove-record-on-delete" - Whether to remove the whole row in aggregation /// engine when delete records are received. Default value is "false". static const char AGGREGATION_REMOVE_RECORD_ON_DELETE[]; diff --git a/include/paimon/global_index/global_indexer.h b/include/paimon/global_index/global_indexer.h index 4da6293ff..690ec4bd1 100644 --- a/include/paimon/global_index/global_indexer.h +++ b/include/paimon/global_index/global_indexer.h @@ -35,6 +35,8 @@ struct ArrowSchema; namespace paimon { +class Executor; + /// Interface for creating global index readers and writers. class PAIMON_EXPORT GlobalIndexer { public: @@ -70,6 +72,27 @@ class PAIMON_EXPORT GlobalIndexer { ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, const std::vector& files, const std::shared_ptr& pool) const = 0; + + /// Creates a reader using an executor supplied by the scan layer. + /// + /// Index implementations which do not perform asynchronous work may ignore the executor and + /// use the compatibility overload above. + /// + /// @param arrow_schema Schema of the indexed data; used to interpret predicate literals. + /// @param file_reader I/O handler for reading index artifacts from storage. + /// @param files List of index file metadata entries produced during writing. + /// @param pool Memory pool for temporary allocations; if nullptr, uses default. + /// @param executor Executor shared by readers created for the same scan; nullptr means + /// that the reader should evaluate sequentially. + /// @return A `Result` containing a shared pointer to the created `GlobalIndexReader`, + /// or an error if the index cannot be loaded or is incompatible, etc. + virtual Result> CreateReader( + ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, + const std::vector& files, const std::shared_ptr& pool, + const std::shared_ptr& executor) const { + static_cast(executor); + return CreateReader(arrow_schema, file_reader, files, pool); + } }; } // namespace paimon diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 514c1354d..25d86c0ef 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -254,6 +254,11 @@ set(PAIMON_CORE_SRCS core/index/index_file_handler.cpp core/index/global_index_meta.cpp core/index/index_file_meta_serializer.cpp + core/index/pk/primary_key_index_source_meta.cpp + core/index/pk/primary_key_index_definitions.cpp + core/index/pksorted/pk_sorted_index_group.cpp + core/index/pksorted/pk_sorted_bucket_index_state.cpp + core/index/pksorted/pk_sorted_index_file.cpp core/io/generic_row_to_arrow_array_converter.cpp core/io/meta_to_arrow_array_converter.cpp core/io/async_key_value_producer_and_consumer.cpp @@ -404,6 +409,9 @@ set(PAIMON_CORE_SRCS core/table/source/table_read.cpp core/table/source/table_scan.cpp core/table/source/data_evolution_batch_scan.cpp + core/table/source/primary_key_sorted_index_scan.cpp + core/table/source/primary_key_sorted_index_result.cpp + core/table/source/primary_key_index_batch_scan.cpp core/table/system/audit_log_system_table.cpp core/table/system/binlog_system_table.cpp core/table/system/global_system_tables.cpp @@ -722,6 +730,9 @@ if(PAIMON_BUILD_TESTS) core/index/index_in_data_file_dir_path_factory_test.cpp core/index/deletion_vector_meta_test.cpp core/index/index_file_meta_serializer_test.cpp + core/index/pk/primary_key_index_source_meta_test.cpp + core/index/pk/primary_key_index_definitions_test.cpp + core/index/pksorted/pk_sorted_bucket_index_state_test.cpp core/index/index_file_handler_test.cpp core/io/compact_increment_test.cpp core/io/infer_shredding_file_writer_test.cpp @@ -857,6 +868,7 @@ if(PAIMON_BUILD_TESTS) core/table/sink/commit_message_test.cpp core/table/sink/commit_message_impl_test.cpp core/table/source/fallback_data_split_test.cpp + core/table/source/primary_key_sorted_index_scan_test.cpp core/table/source/table_read_test.cpp core/table/source/append_count_reader_test.cpp core/table/source/pk_count_reader_test.cpp diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index 53f18c3fd..58d8b4bb6 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -139,6 +139,10 @@ const char Options::BLOB_WRITE_NULL_ON_FETCH_FAILURE[] = "blob-write-null-on-fet const char Options::GLOBAL_INDEX_ENABLED[] = "global-index.enabled"; const char Options::GLOBAL_INDEX_THREAD_NUM[] = "global-index.thread-num"; const char Options::GLOBAL_INDEX_EXTERNAL_PATH[] = "global-index.external-path"; +const char Options::PK_BTREE_INDEX_COLUMNS[] = "pk-btree.index.columns"; +const char Options::PK_BITMAP_INDEX_COLUMNS[] = "pk-bitmap.index.columns"; +const char Options::PK_VECTOR_INDEX_COLUMNS[] = "pk-vector.index.columns"; +const char Options::PK_FULL_TEXT_INDEX_COLUMNS[] = "pk-full-text.index.columns"; const char Options::AGGREGATION_REMOVE_RECORD_ON_DELETE[] = "aggregation.remove-record-on-delete"; const char Options::TABLE_READ_SEQUENCE_NUMBER_ENABLED[] = "table-read.sequence-number.enabled"; const char Options::KEY_VALUE_SEQUENCE_NUMBER_ENABLED[] = "key-value.sequence_number.enabled"; 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 e5aeaac20..9bfb97e0c 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 @@ -16,6 +16,10 @@ * specific language governing permissions and limitations * under the License. */ +#include +#include +#include + #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" @@ -29,6 +33,7 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/data/decimal.h" #include "paimon/data/timestamp.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" @@ -84,6 +89,27 @@ class FakeGlobalIndexFileReader : public GlobalIndexFileReader { std::string base_path_; }; +class CountingInlineExecutor : public Executor { + public: + void Add(std::function func) override { + submission_count_.fetch_add(1); + func(); + } + + void ShutdownNow() override {} + + uint32_t GetThreadNum() const override { + return 1; + } + + uint32_t SubmissionCount() const { + return submission_count_.load(); + } + + private: + std::atomic submission_count_{0}; +}; + class BTreeGlobalIndexIntegrationTest : public ::testing::Test, public ::testing::WithParamInterface { protected: @@ -1972,9 +1998,10 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadMultiFilesWithMetaSelector) // Create reader over all 3 files (internally uses LazyFilteredBTreeReader + // BTreeFileMetaSelector) auto file_reader = std::make_shared(fs_, base_path_); + auto executor = std::make_shared(); auto c_schema = CreateArrowSchema(field); - ASSERT_OK_AND_ASSIGN(auto reader, - indexer->CreateReader(c_schema.get(), file_reader, all_metas, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, indexer->CreateReader(c_schema.get(), file_reader, all_metas, + pool_, executor)); // --- VisitEqual: key=12 -> only file1 is selected by meta selector -> row 5 { @@ -2023,6 +2050,7 @@ TEST_P(BTreeGlobalIndexIntegrationTest, WriteAndReadMultiFilesWithMetaSelector) Literal literal_5(5); ASSERT_OK_AND_ASSIGN(auto result, reader->VisitGreaterOrEqual(literal_5)); CheckResult(result, {3, 4, 5, 6, 7, 9}); + ASSERT_EQ(executor->SubmissionCount(), 3); } // --- VisitLessOrEqual: key <= 2 -> only file0 selected -> rows 0,1 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 29bffe9fe..062b27e61 100644 --- a/src/paimon/common/global_index/btree/btree_global_indexer.cpp +++ b/src/paimon/common/global_index/btree/btree_global_indexer.cpp @@ -18,10 +18,12 @@ */ #include "paimon/common/global_index/btree/btree_global_indexer.h" +#include #include #include #include "arrow/c/bridge.h" +#include "fmt/format.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_writer.h" @@ -36,10 +38,10 @@ #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" + namespace paimon { Result> BTreeGlobalIndexer::Create( const std::map& options) { @@ -100,6 +102,13 @@ Result> BTreeGlobalIndexer::CreateWriter( Result> BTreeGlobalIndexer::CreateReader( ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, const std::vector& files, const std::shared_ptr& pool) const { + return CreateReader(arrow_schema, file_reader, files, pool, /*executor=*/nullptr); +} + +Result> BTreeGlobalIndexer::CreateReader( + ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, + const std::vector& files, const std::shared_ptr& pool, + const std::shared_ptr& executor) const { // Get field type from arrow schema PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, arrow::ImportSchema(arrow_schema)); @@ -120,8 +129,6 @@ Result> BTreeGlobalIndexer::CreateReader( } read_buffer_size = static_cast(tmp_buffer_size); } - // TODO(lisizhuo.lsz): Allow users to specify an executor - std::shared_ptr executor = CreateDefaultExecutor(); return std::make_shared(read_buffer_size, files, key_type, file_reader, cache_manager_, pool, executor); } 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 5568adba3..f93099487 100644 --- a/src/paimon/common/global_index/btree/btree_global_indexer.h +++ b/src/paimon/common/global_index/btree/btree_global_indexer.h @@ -68,6 +68,11 @@ class BTreeGlobalIndexer : public GlobalIndexer { const std::vector& files, const std::shared_ptr& pool) const override; + Result> CreateReader( + ::ArrowSchema* arrow_schema, const std::shared_ptr& file_reader, + const std::vector& files, const std::shared_ptr& pool, + const std::shared_ptr& executor) const override; + private: BTreeGlobalIndexer(const std::shared_ptr& cache_manager, const std::map& options) diff --git a/src/paimon/core/index/pk/primary_key_index_definition.h b/src/paimon/core/index/pk/primary_key_index_definition.h new file mode 100644 index 000000000..1a06628e9 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_definition.h @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include + +namespace paimon { +/// Resolved definition of one source-backed primary-key index. +class PrimaryKeyIndexDefinition { + public: + /// Built-in primary-key index families. + enum class Family { + VECTOR, + BTREE, + BITMAP, + FULL_TEXT, + }; + + PrimaryKeyIndexDefinition(std::string column, int32_t field_id, std::string index_type, + Family family, std::map options) + : column_(std::move(column)), + field_id_(field_id), + index_type_(std::move(index_type)), + family_(family), + options_(std::move(options)) {} + + const std::string& Column() const { + return column_; + } + + int32_t FieldId() const { + return field_id_; + } + + const std::string& IndexType() const { + return index_type_; + } + + Family GetFamily() const { + return family_; + } + + const std::map& Options() const { + return options_; + } + + private: + std::string column_; + int32_t field_id_; + std::string index_type_; + Family family_; + std::map options_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_definitions.cpp b/src/paimon/core/index/pk/primary_key_index_definitions.cpp new file mode 100644 index 000000000..8f87184ba --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_definitions.cpp @@ -0,0 +1,223 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pk/primary_key_index_definitions.h" + +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/utils/object_utils.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/defs.h" +#include "rapidjson/document.h" +#include "rapidjson/stringbuffer.h" +#include "rapidjson/writer.h" + +namespace paimon { +namespace { +using IndexOptions = std::map; + +constexpr char kBTreeIndexType[] = "btree"; +constexpr char kBitmapIndexType[] = "bitmap"; +constexpr char kFullTextIndexType[] = "full-text"; +constexpr char kBTreeOptionFamily[] = "pk-btree"; +constexpr char kBitmapOptionFamily[] = "pk-bitmap"; +constexpr char kBTreeAlgorithmPrefix[] = "btree-index."; +constexpr char kBitmapAlgorithmPrefix[] = "bitmap-index."; +constexpr char kFieldScopedPrefix[] = "fields."; +constexpr char kRecordsPerRangeKey[] = "sorted-index.records-per-range"; + +std::vector IndexColumns(const std::map& options, + const char* option_key) { + auto iter = options.find(option_key); + if (iter == options.end()) { + return {}; + } + std::vector columns = StringUtils::Split(iter->second, ",", false); + for (std::string& column : columns) { + StringUtils::Trim(&column); + } + return columns; +} + +Status AddUniqueColumns(const std::vector& columns, + const std::function& on_duplicate, + std::set* unique_columns) { + for (const std::string& column : columns) { + if (!unique_columns->insert(column).second) { + return on_duplicate(column); + } + } + return Status::OK(); +} + +Status ValidateNoDuplicates(const std::vector& columns, const char* option_key) { + std::set unique_columns; + return AddUniqueColumns( + columns, + [option_key](const std::string& column) { + return Status::Invalid( + fmt::format("{} contains duplicate column '{}'.", option_key, column)); + }, + &unique_columns); +} + +Status ValidateUniqueColumns(const std::vector& columns, + std::set* indexed_columns) { + return AddUniqueColumns( + columns, + [](const std::string& column) { + return Status::Invalid( + fmt::format("Column '{}' can own at most one primary-key index.", column)); + }, + indexed_columns); +} + +/// Resolves the effective option map of one sorted-index definition: table options first, +/// then the field-scoped JSON options with unqualified keys prefixed by the algorithm +/// prefix, mirroring Java `CoreOptions#primaryKeySortedIndexOptions`. +Result> SortedIndexOptions( + const std::map& table_options, const std::string& column, + const char* option_family, const char* algorithm_prefix) { + std::map resolved = table_options; + resolved.erase(kRecordsPerRangeKey); + std::string option_key = + fmt::format("{}{}.{}.index.options", kFieldScopedPrefix, column, option_family); + auto iter = table_options.find(option_key); + if (iter == table_options.end() || StringUtils::IsNullOrWhitespaceOnly(iter->second)) { + return resolved; + } + + rapidjson::Document document; + document.Parse(iter->second.c_str()); + if (document.HasParseError() || !document.IsObject()) { + return Status::Invalid( + fmt::format("{} must be a JSON object of option key-value pairs.", option_key)); + } + for (auto member = document.MemberBegin(); member != document.MemberEnd(); ++member) { + if (!member->name.IsString() || + StringUtils::IsNullOrWhitespaceOnly(member->name.GetString())) { + return Status::Invalid(fmt::format("{} contains an empty option key.", option_key)); + } + std::string key = member->name.GetString(); + if (member->value.IsNull()) { + return Status::Invalid( + fmt::format("{} value for key {} must not be null.", option_key, key)); + } + if (member->value.IsObject() || member->value.IsArray()) { + return Status::Invalid( + fmt::format("{} must be a JSON object of option key-value pairs.", option_key)); + } + std::string value; + if (member->value.IsString()) { + value = member->value.GetString(); + } else { + // Java's parseJsonMap(..., String.class) coerces scalar JSON values (numbers, + // booleans) to their text form, so `{"compression-level":3}` is valid there. + rapidjson::StringBuffer buffer; + rapidjson::Writer writer(buffer); + member->value.Accept(writer); + value = buffer.GetString(); + } + std::string qualified_key = StringUtils::StartsWith(key, algorithm_prefix) || + StringUtils::StartsWith(key, kFieldScopedPrefix) + ? key + : algorithm_prefix + key; + auto previous = resolved.find(qualified_key); + if (previous != resolved.end() && previous->second != value) { + return Status::Invalid( + fmt::format("{} defines conflicting values for {}.", option_key, qualified_key)); + } + resolved[qualified_key] = value; + } + return resolved; +} + +} // namespace + +Result PrimaryKeyIndexDefinitions::Create(const TableSchema& schema) { + const std::map& options = schema.Options(); + std::vector vector_columns = + IndexColumns(options, Options::PK_VECTOR_INDEX_COLUMNS); + std::vector btree_columns = IndexColumns(options, Options::PK_BTREE_INDEX_COLUMNS); + std::vector bitmap_columns = + IndexColumns(options, Options::PK_BITMAP_INDEX_COLUMNS); + std::vector full_text_columns = + IndexColumns(options, Options::PK_FULL_TEXT_INDEX_COLUMNS); + PAIMON_RETURN_NOT_OK(ValidateNoDuplicates(vector_columns, Options::PK_VECTOR_INDEX_COLUMNS)); + PAIMON_RETURN_NOT_OK(ValidateNoDuplicates(btree_columns, Options::PK_BTREE_INDEX_COLUMNS)); + PAIMON_RETURN_NOT_OK(ValidateNoDuplicates(bitmap_columns, Options::PK_BITMAP_INDEX_COLUMNS)); + PAIMON_RETURN_NOT_OK( + ValidateNoDuplicates(full_text_columns, Options::PK_FULL_TEXT_INDEX_COLUMNS)); + std::set indexed_columns; + PAIMON_RETURN_NOT_OK(ValidateUniqueColumns(vector_columns, &indexed_columns)); + PAIMON_RETURN_NOT_OK(ValidateUniqueColumns(btree_columns, &indexed_columns)); + PAIMON_RETURN_NOT_OK(ValidateUniqueColumns(bitmap_columns, &indexed_columns)); + PAIMON_RETURN_NOT_OK(ValidateUniqueColumns(full_text_columns, &indexed_columns)); + + std::vector definitions; + for (const DataField& field : schema.Fields()) { + const std::string& column = field.Name(); + if (ObjectUtils::Contains(btree_columns, column)) { + PAIMON_ASSIGN_OR_RAISE( + IndexOptions definition_options, + SortedIndexOptions(options, column, kBTreeOptionFamily, kBTreeAlgorithmPrefix)); + definitions.emplace_back(column, field.Id(), kBTreeIndexType, + PrimaryKeyIndexDefinition::Family::BTREE, + std::move(definition_options)); + } else if (ObjectUtils::Contains(bitmap_columns, column)) { + PAIMON_ASSIGN_OR_RAISE( + IndexOptions definition_options, + SortedIndexOptions(options, column, kBitmapOptionFamily, kBitmapAlgorithmPrefix)); + definitions.emplace_back(column, field.Id(), kBitmapIndexType, + PrimaryKeyIndexDefinition::Family::BITMAP, + std::move(definition_options)); + } else if (ObjectUtils::Contains(vector_columns, column)) { + std::string index_type; + auto type_iter = + options.find(fmt::format("{}{}.pk-vector.index.type", kFieldScopedPrefix, column)); + if (type_iter != options.end()) { + index_type = type_iter->second; + } + definitions.emplace_back(column, field.Id(), index_type, + PrimaryKeyIndexDefinition::Family::VECTOR, + std::map()); + } else if (ObjectUtils::Contains(full_text_columns, column)) { + definitions.emplace_back(column, field.Id(), kFullTextIndexType, + PrimaryKeyIndexDefinition::Family::FULL_TEXT, + std::map()); + } + } + return PrimaryKeyIndexDefinitions(std::move(definitions)); +} + +std::vector PrimaryKeyIndexDefinitions::ScalarDefinitions() const { + std::vector scalar_definitions; + for (const PrimaryKeyIndexDefinition& definition : definitions_) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE || + definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BITMAP) { + scalar_definitions.push_back(definition); + } + } + return scalar_definitions; +} + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_definitions.h b/src/paimon/core/index/pk/primary_key_index_definitions.h new file mode 100644 index 000000000..37f20da5f --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_definitions.h @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/result.h" + +namespace paimon { +/// Resolves all configured source-backed primary-key indexes of a table schema. +class PrimaryKeyIndexDefinitions { + public: + /// Resolves index definitions from `pk-btree.index.columns`, `pk-bitmap.index.columns`, + /// `pk-vector.index.columns` and `pk-full-text.index.columns` together with their + /// field-scoped option JSON, rejecting duplicate columns and columns owned by more than + /// one index family. + static Result Create(const TableSchema& schema); + + const std::vector& Definitions() const { + return definitions_; + } + + /// @return The scalar (BTree / Bitmap) definitions usable by batch scans. + std::vector ScalarDefinitions() const; + + private: + explicit PrimaryKeyIndexDefinitions(std::vector definitions) + : definitions_(std::move(definitions)) {} + + std::vector definitions_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_definitions_test.cpp b/src/paimon/core/index/pk/primary_key_index_definitions_test.cpp new file mode 100644 index 000000000..616b0dc98 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_definitions_test.cpp @@ -0,0 +1,216 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pk/primary_key_index_definitions.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "gtest/gtest.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/defs.h" +#include "paimon/result.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { +/// Builds a primary-key table schema with fields id BIGINT (pk), price DOUBLE, age INT, +/// status STRING and emb FLOAT, merging the given options over a fixed bucket option. +Result> MakeSchema(std::map options) { + std::vector fields = { + DataField(0, arrow::field("id", arrow::int64(), /*nullable=*/false)), + DataField(1, arrow::field("price", arrow::float64())), + DataField(2, arrow::field("age", arrow::int32())), + DataField(3, arrow::field("status", arrow::utf8())), + DataField(4, arrow::field("emb", arrow::float32()))}; + options.emplace(Options::BUCKET, "1"); + return TableSchema::Create(/*schema_id=*/0, DataField::ConvertDataFieldsToArrowSchema(fields), + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options); +} +} // namespace + +TEST(PrimaryKeyIndexDefinitionsTest, NoIndexOptionsYieldsEmptyDefinitions) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, MakeSchema({})); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_TRUE(definitions.Definitions().empty()); + ASSERT_TRUE(definitions.ScalarDefinitions().empty()); +} + +TEST(PrimaryKeyIndexDefinitionsTest, ResolvesBTreeDefinitions) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price,age"}})); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_EQ(2, definitions.Definitions().size()); + const PrimaryKeyIndexDefinition& price = definitions.Definitions()[0]; + ASSERT_EQ("price", price.Column()); + ASSERT_EQ(1, price.FieldId()); + ASSERT_EQ("btree", price.IndexType()); + ASSERT_EQ(PrimaryKeyIndexDefinition::Family::BTREE, price.GetFamily()); + const PrimaryKeyIndexDefinition& age = definitions.Definitions()[1]; + ASSERT_EQ("age", age.Column()); + ASSERT_EQ(2, age.FieldId()); + ASSERT_EQ("btree", age.IndexType()); + ASSERT_EQ(PrimaryKeyIndexDefinition::Family::BTREE, age.GetFamily()); + ASSERT_EQ(2, definitions.ScalarDefinitions().size()); +} + +TEST(PrimaryKeyIndexDefinitionsTest, ResolvesBitmapDefinition) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BITMAP_INDEX_COLUMNS, "status"}})); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_EQ(1, definitions.Definitions().size()); + const PrimaryKeyIndexDefinition& status = definitions.Definitions()[0]; + ASSERT_EQ("status", status.Column()); + ASSERT_EQ(3, status.FieldId()); + ASSERT_EQ("bitmap", status.IndexType()); + ASSERT_EQ(PrimaryKeyIndexDefinition::Family::BITMAP, status.GetFamily()); + ASSERT_EQ(1, definitions.ScalarDefinitions().size()); +} + +TEST(PrimaryKeyIndexDefinitionsTest, IgnoresColumnAbsentFromSchema) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "not_in_schema"}})); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_TRUE(definitions.Definitions().empty()); +} + +TEST(PrimaryKeyIndexDefinitionsTest, CoercesScalarJsonOptionValuesLikeJava) { + // Java's parseJsonMap(..., String.class) accepts scalar JSON values and coerces them + // to text, so numeric or boolean values written by a Java engine must stay readable. + std::map options = { + {Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"fields.price.pk-btree.index.options", + R"({"compression-level":3,"cache-enabled":true,"block-size":"64 kb"})"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, MakeSchema(options)); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + const std::map& resolved = definitions.Definitions()[0].Options(); + ASSERT_EQ("3", resolved.at("btree-index.compression-level")); + ASSERT_EQ("true", resolved.at("btree-index.cache-enabled")); + ASSERT_EQ("64 kb", resolved.at("btree-index.block-size")); + + // Null and nested values are rejected like in Java. + std::map null_options = { + {Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"fields.price.pk-btree.index.options", R"({"block-size":null})"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr null_schema, MakeSchema(null_options)); + ASSERT_NOK(PrimaryKeyIndexDefinitions::Create(*null_schema)); + std::map nested_options = { + {Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"fields.price.pk-btree.index.options", R"({"block-size":{"v":"64 kb"}})"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr nested_schema, MakeSchema(nested_options)); + ASSERT_NOK(PrimaryKeyIndexDefinitions::Create(*nested_schema)); +} + +TEST(PrimaryKeyIndexDefinitionsTest, QualifiesFieldScopedJsonOptions) { + std::map options = { + {Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"sorted-index.records-per-range", "4096"}, + {"fields.price.pk-btree.index.options", + R"({"block-size":"64 kb","btree-index.cache-size":"32 mb","fields.foo.x":"y"})"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, MakeSchema(options)); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_EQ(1, definitions.Definitions().size()); + const std::map& resolved = definitions.Definitions()[0].Options(); + // Unqualified keys are prefixed with the algorithm prefix, qualified keys are kept as-is. + ASSERT_EQ(1, resolved.count("btree-index.block-size")); + ASSERT_EQ("64 kb", resolved.at("btree-index.block-size")); + ASSERT_EQ(1, resolved.count("btree-index.cache-size")); + ASSERT_EQ("32 mb", resolved.at("btree-index.cache-size")); + ASSERT_EQ(1, resolved.count("fields.foo.x")); + ASSERT_EQ("y", resolved.at("fields.foo.x")); + // The per-range knob never leaks into the definition, other table options are retained. + ASSERT_EQ(0, resolved.count("sorted-index.records-per-range")); + ASSERT_EQ(1, resolved.count(Options::BUCKET)); + ASSERT_EQ("1", resolved.at(Options::BUCKET)); +} + +TEST(PrimaryKeyIndexDefinitionsTest, RejectsConflictingJsonOptionValue) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"btree-index.block-size", "128 kb"}, + {"fields.price.pk-btree.index.options", R"({"block-size":"64 kb"})"}})); + ASSERT_NOK(PrimaryKeyIndexDefinitions::Create(*schema)); +} + +TEST(PrimaryKeyIndexDefinitionsTest, RejectsMalformedJsonOptions) { + { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"fields.price.pk-btree.index.options", "not-json"}})); + ASSERT_NOK(PrimaryKeyIndexDefinitions::Create(*schema)); + } + { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {"fields.price.pk-btree.index.options", R"({"":"v"})"}})); + ASSERT_NOK(PrimaryKeyIndexDefinitions::Create(*schema)); + } +} + +TEST(PrimaryKeyIndexDefinitionsTest, RejectsDuplicateColumnWithinFamily) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price,price"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyIndexDefinitions::Create(*schema), + "pk-btree.index.columns contains duplicate column 'price'."); +} + +TEST(PrimaryKeyIndexDefinitionsTest, RejectsColumnSharedAcrossFamilies) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {Options::PK_BITMAP_INDEX_COLUMNS, "price"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyIndexDefinitions::Create(*schema), + "Column 'price' can own at most one primary-key index."); +} + +TEST(PrimaryKeyIndexDefinitionsTest, ResolvesNonScalarFamiliesAndExcludesThemFromScalar) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakeSchema({{Options::PK_BTREE_INDEX_COLUMNS, "price"}, + {Options::PK_VECTOR_INDEX_COLUMNS, "emb"}, + {"fields.emb.pk-vector.index.type", "ivf-flat"}, + {Options::PK_FULL_TEXT_INDEX_COLUMNS, "status"}})); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + ASSERT_EQ(3, definitions.Definitions().size()); + const PrimaryKeyIndexDefinition& full_text = definitions.Definitions()[1]; + ASSERT_EQ("status", full_text.Column()); + ASSERT_EQ("full-text", full_text.IndexType()); + ASSERT_EQ(PrimaryKeyIndexDefinition::Family::FULL_TEXT, full_text.GetFamily()); + const PrimaryKeyIndexDefinition& embedding = definitions.Definitions()[2]; + ASSERT_EQ("emb", embedding.Column()); + ASSERT_EQ(4, embedding.FieldId()); + ASSERT_EQ("ivf-flat", embedding.IndexType()); + ASSERT_EQ(PrimaryKeyIndexDefinition::Family::VECTOR, embedding.GetFamily()); + std::vector scalar_definitions = definitions.ScalarDefinitions(); + ASSERT_EQ(1, scalar_definitions.size()); + ASSERT_EQ("price", scalar_definitions[0].Column()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/index/pk/primary_key_index_source_file.h b/src/paimon/core/index/pk/primary_key_index_source_file.h new file mode 100644 index 000000000..c1afa3de7 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_source_file.h @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include + +namespace paimon { +/// One ordered source data file covered by a source-backed primary-key index payload. +/// +/// Two source files are interchangeable only when both the file name and the row count +/// match; coverage validation relies on this strict identity. +struct PrimaryKeyIndexSourceFile { + PrimaryKeyIndexSourceFile(std::string file_name, int64_t row_count) + : file_name(std::move(file_name)), row_count(row_count) {} + + bool operator==(const PrimaryKeyIndexSourceFile& other) const { + return file_name == other.file_name && row_count == other.row_count; + } + + bool operator!=(const PrimaryKeyIndexSourceFile& other) const { + return !(*this == other); + } + + std::string file_name; + int64_t row_count; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_source_meta.cpp b/src/paimon/core/index/pk/primary_key_index_source_meta.cpp new file mode 100644 index 000000000..c31c34689 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_source_meta.cpp @@ -0,0 +1,161 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pk/primary_key_index_source_meta.h" + +#include +#include +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/memory/memory_segment_utils.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/io/data_input_stream.h" + +namespace paimon { +namespace { +// Each serialized entry needs at least one uint16 string length and one int64 row count, +// mirroring the defensive source file count cap of the Java deserializer. +constexpr int64_t kMinBytesPerSourceFile = sizeof(uint16_t) + sizeof(int64_t); +constexpr size_t kMaxInitialSourceFileCapacity = 1024; +} // namespace + +Result PrimaryKeyIndexSourceMeta::Create( + int32_t data_level, std::vector source_files) { + if (data_level <= 0) { + return Status::Invalid("Primary-key index data level must be positive."); + } + if (source_files.empty()) { + return Status::Invalid("An index must reference source files."); + } + for (const PrimaryKeyIndexSourceFile& source_file : source_files) { + if (source_file.row_count < 0) { + return Status::Invalid(fmt::format("Source file {} has a negative row count {}.", + source_file.file_name, source_file.row_count)); + } + } + return PrimaryKeyIndexSourceMeta(data_level, std::move(source_files)); +} + +Result PrimaryKeyIndexSourceMeta::FromIndexFile( + const IndexFileMeta& index_file) { + const std::optional& global_index_meta = index_file.GetGlobalIndexMeta(); + if (global_index_meta == std::nullopt || global_index_meta.value().source_meta == nullptr) { + return Status::Invalid( + fmt::format("Index file {} has no source metadata.", index_file.FileName())); + } + const std::shared_ptr& source_meta = global_index_meta.value().source_meta; + return Deserialize(source_meta->data(), source_meta->size()); +} + +Result PrimaryKeyIndexSourceMeta::Deserialize(const char* data, + size_t length) { + if (data == nullptr) { + return Status::Invalid("Cannot deserialize index source metadata from a null buffer."); + } + if (length > static_cast(std::numeric_limits::max())) { + return Status::Invalid( + fmt::format("Index source metadata length {} exceeds the supported maximum {}.", length, + std::numeric_limits::max())); + } + + auto input_stream = std::make_shared(data, static_cast(length)); + DataInputStream input(input_stream); + PAIMON_ASSIGN_OR_RAISE(int32_t version, input.ReadValue()); + if (version != VERSION) { + return Status::Invalid(fmt::format("Unsupported index source version: {}.", version)); + } + PAIMON_ASSIGN_OR_RAISE(int32_t data_level, input.ReadValue()); + PAIMON_ASSIGN_OR_RAISE(int32_t source_file_count, input.ReadValue()); + if (source_file_count <= 0) { + return Status::Invalid("An index must reference source files."); + } + PAIMON_ASSIGN_OR_RAISE(int64_t position, input.GetPos()); + PAIMON_ASSIGN_OR_RAISE(int64_t stream_length, input.Length()); + int64_t maximum_source_file_count = (stream_length - position) / kMinBytesPerSourceFile; + if (static_cast(source_file_count) > maximum_source_file_count) { + return Status::Invalid(fmt::format( + "Failed to deserialize index source metadata: source file count {} exceeds the " + "maximum {} allowed by the remaining bytes.", + source_file_count, maximum_source_file_count)); + } + std::vector source_files; + source_files.reserve( + std::min(static_cast(source_file_count), kMaxInitialSourceFileCapacity)); + for (int32_t i = 0; i < source_file_count; i++) { + PAIMON_ASSIGN_OR_RAISE(std::string file_name, input.ReadString()); + PAIMON_ASSIGN_OR_RAISE(int64_t row_count, input.ReadValue()); + source_files.emplace_back(std::move(file_name), row_count); + } + PAIMON_ASSIGN_OR_RAISE(position, input.GetPos()); + if (position != stream_length) { + return Status::Invalid("Unexpected trailing bytes in index source metadata."); + } + return Create(data_level, std::move(source_files)); +} + +Result> PrimaryKeyIndexSourceMeta::Serialize( + const std::shared_ptr& pool) const { + if (pool == nullptr) { + return Status::Invalid("Cannot serialize index source metadata with a null memory pool."); + } + if (source_files_.size() > static_cast(std::numeric_limits::max())) { + return Status::Invalid( + fmt::format("Index source file count {} exceeds the supported maximum {}.", + source_files_.size(), std::numeric_limits::max())); + } + + int64_t serialized_size = 3 * static_cast(sizeof(int32_t)); + for (const PrimaryKeyIndexSourceFile& source_file : source_files_) { + if (source_file.file_name.size() > std::numeric_limits::max()) { + return Status::Invalid( + fmt::format("Source file name is too long for a 16-bit length: {} bytes.", + source_file.file_name.size())); + } + int64_t entry_size = + kMinBytesPerSourceFile + static_cast(source_file.file_name.size()); + if (serialized_size > std::numeric_limits::max() - entry_size) { + return Status::Invalid(fmt::format( + "Serialized index source metadata exceeds the supported maximum {} bytes.", + std::numeric_limits::max())); + } + serialized_size += entry_size; + } + + MemorySegmentOutputStream output(MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool); + output.WriteValue(VERSION); + output.WriteValue(data_level_); + output.WriteValue(static_cast(source_files_.size())); + for (const PrimaryKeyIndexSourceFile& source_file : source_files_) { + auto name_length = static_cast(source_file.file_name.size()); + output.WriteValue(name_length); + output.Write(source_file.file_name.data(), name_length); + output.WriteValue(source_file.row_count); + } + PAIMON_UNIQUE_PTR bytes = MemorySegmentUtils::CopyToBytes( + output.Segments(), /*offset=*/0, static_cast(serialized_size), pool.get()); + return std::shared_ptr(std::move(bytes)); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_source_meta.h b/src/paimon/core/index/pk/primary_key_index_source_meta.h new file mode 100644 index 000000000..dbfa485b2 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_source_meta.h @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace paimon { +class IndexFileMeta; + +/// Ordered source data files covered by a source-backed primary-key index payload. +/// +/// Wire format (version 1): big-endian int32 version, big-endian int32 data level (> 0), +/// big-endian int32 source file count (> 0), then per source file a uint16 big-endian byte +/// length, unchanged file name bytes, and a big-endian int64 row count. This matches Java +/// `writeUTF` for ASCII and non-null BMP UTF-8 file names. Java modified UTF-8 support for +/// supplementary code points requires a stream-level follow-up. Trailing bytes are rejected. +class PrimaryKeyIndexSourceMeta { + public: + static constexpr int32_t VERSION = 1; + + static Result Create( + int32_t data_level, std::vector source_files); + + /// Extracts and deserializes the source metadata carried by an index file. + static Result FromIndexFile(const IndexFileMeta& index_file); + + static Result Deserialize(const char* data, size_t length); + + Result> Serialize(const std::shared_ptr& pool) const; + + int32_t DataLevel() const { + return data_level_; + } + + const std::vector& SourceFiles() const { + return source_files_; + } + + private: + PrimaryKeyIndexSourceMeta(int32_t data_level, + std::vector source_files) + : data_level_(data_level), source_files_(std::move(source_files)) {} + + int32_t data_level_; + std::vector source_files_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pk/primary_key_index_source_meta_test.cpp b/src/paimon/core/index/pk/primary_key_index_source_meta_test.cpp new file mode 100644 index 000000000..44b83f5d3 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_source_meta_test.cpp @@ -0,0 +1,223 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pk/primary_key_index_source_meta.h" + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class PrimaryKeyIndexSourceMetaTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + } + + Result SerializeToString(int32_t data_level, + std::vector source_files) { + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeyIndexSourceMeta meta, + PrimaryKeyIndexSourceMeta::Create(data_level, std::move(source_files))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, meta.Serialize(pool_)); + return std::string(bytes->data(), bytes->size()); + } + + std::shared_ptr pool_; +}; + +TEST_F(PrimaryKeyIndexSourceMetaTest, SerializeMatchesGoldenBytes) { + std::vector files; + files.emplace_back("a.parquet", 100); + files.emplace_back("b.parquet", 200); + ASSERT_OK_AND_ASSIGN(std::string serialized, SerializeToString(3, files)); + + const uint8_t kExpected[] = { + 0x00, 0x00, 0x00, 0x01, // version 1 + 0x00, 0x00, 0x00, 0x03, // data level 3 + 0x00, 0x00, 0x00, 0x02, // source file count 2 + 0x00, 0x09, // writeUTF byte length of "a.parquet" + 'a', '.', 'p', 'a', 'r', 'q', 'u', 'e', 't', + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x64, // row count 100 + 0x00, 0x09, // writeUTF byte length of "b.parquet" + 'b', '.', 'p', 'a', 'r', 'q', 'u', 'e', 't', + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xC8, // row count 200 + }; + std::string expected(reinterpret_cast(kExpected), sizeof(kExpected)); + ASSERT_EQ(expected, serialized); + + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta meta, PrimaryKeyIndexSourceMeta::Deserialize( + serialized.data(), serialized.size())); + ASSERT_EQ(3, meta.DataLevel()); + ASSERT_EQ(files, meta.SourceFiles()); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, RoundTripWithBmpNameAndLargeRowCount) { + std::vector files; + // Java modified UTF-8 and standard UTF-8 use identical bytes for non-null BMP text. + // Row count above 2^32 exercises the full big-endian int64 encoding. + files.emplace_back("文件-0.parquet", (int64_t{1} << 40) + 7); + files.emplace_back("data-1.parquet", 42); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta meta, + PrimaryKeyIndexSourceMeta::Create(5, files)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr bytes, meta.Serialize(pool_)); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta decoded, + PrimaryKeyIndexSourceMeta::Deserialize(bytes->data(), bytes->size())); + ASSERT_EQ(5, decoded.DataLevel()); + ASSERT_EQ(files, decoded.SourceFiles()); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, DeserializeRejectsBadHeaders) { + std::vector files; + files.emplace_back("a.parquet", 100); + ASSERT_OK_AND_ASSIGN(std::string valid, SerializeToString(3, files)); + // Layout: [0,4) version, [4,8) data level, [8,12) count, [12,14) name length, + // [14,23) name bytes, [23,31) row count. + ASSERT_EQ(static_cast(31), valid.size()); + + // Unsupported versions. + std::string version_two = valid; + version_two[3] = '\x02'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(version_two.data(), version_two.size())); + std::string version_zero = valid; + version_zero[3] = '\x00'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(version_zero.data(), version_zero.size())); + + // Source file count must be positive. + std::string zero_count = valid; + zero_count[11] = '\x00'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(zero_count.data(), zero_count.size())); + std::string negative_count = valid; + for (size_t i = 8; i < 12; i++) { + negative_count[i] = '\xFF'; + } + ASSERT_NOK( + PrimaryKeyIndexSourceMeta::Deserialize(negative_count.data(), negative_count.size())); + + // Claimed count 1000 exceeds the defensive cap allowed by the 19 remaining bytes. + std::string huge_count = valid; + huge_count[10] = '\x03'; + huge_count[11] = '\xE8'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(huge_count.data(), huge_count.size())); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, DeserializeRejectsBadPayloads) { + std::vector files; + files.emplace_back("a.parquet", 100); + ASSERT_OK_AND_ASSIGN(std::string valid, SerializeToString(3, files)); + ASSERT_EQ(static_cast(31), valid.size()); + + // Trailing bytes after a valid payload. + std::string trailing = valid + '\x00'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(trailing.data(), trailing.size())); + + // Buffer cut in the middle of the file name: only 8 of the 9 name bytes remain. + std::string cut_name = valid.substr(0, 22); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(cut_name.data(), cut_name.size())); + + // Buffer cut in the middle of the row count: only 4 of the 8 bytes remain. + std::string cut_row_count = valid.substr(0, 27); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(cut_row_count.data(), cut_row_count.size())); + + std::string negative_row_count = valid; + negative_row_count[23] = '\xFF'; + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Deserialize(negative_row_count.data(), + negative_row_count.size())); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, CreateRejectsInvalidArguments) { + std::vector files; + files.emplace_back("a.parquet", 100); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Create(0, files)); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Create(-1, files)); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Create(3, {})); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::Create(3, {{"a.parquet", -1}})); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, SerializeValidatesStringLengthAndMemoryPool) { + std::string maximum_name(std::numeric_limits::max(), 'a'); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta maximum_meta, + PrimaryKeyIndexSourceMeta::Create(1, {{maximum_name, 1}})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr maximum_bytes, maximum_meta.Serialize(pool_)); + ASSERT_OK_AND_ASSIGN( + PrimaryKeyIndexSourceMeta maximum_decoded, + PrimaryKeyIndexSourceMeta::Deserialize(maximum_bytes->data(), maximum_bytes->size())); + ASSERT_EQ(maximum_meta.SourceFiles(), maximum_decoded.SourceFiles()); + + std::string oversized_name(static_cast(std::numeric_limits::max()) + 1, 'a'); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta oversized_meta, + PrimaryKeyIndexSourceMeta::Create(1, {{std::move(oversized_name), 1}})); + ASSERT_NOK_WITH_MSG(oversized_meta.Serialize(pool_), "too long for a 16-bit length"); + ASSERT_NOK_WITH_MSG(maximum_meta.Serialize(nullptr), "null memory pool"); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, DeserializeRejectsNullBuffer) { + ASSERT_NOK_WITH_MSG(PrimaryKeyIndexSourceMeta::Deserialize(nullptr, 0), "null buffer"); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, FromIndexFileDecodesSourceMeta) { + std::vector files; + files.emplace_back("a.parquet", 100); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta meta, + PrimaryKeyIndexSourceMeta::Create(7, files)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr source_meta, meta.Serialize(pool_)); + std::shared_ptr index_meta = std::make_shared("index-payload", pool_.get()); + GlobalIndexMeta global_index_meta(/*_row_range_start=*/0, /*_row_range_end=*/100, + /*_index_field_id=*/1, /*_extra_field_ids=*/std::nullopt, + index_meta, source_meta); + IndexFileMeta index_file("pk-btree", "index-file-0", /*file_size=*/64, /*row_count=*/100, + /*dv_ranges=*/std::nullopt, /*external_path=*/std::nullopt, + global_index_meta); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta decoded, + PrimaryKeyIndexSourceMeta::FromIndexFile(index_file)); + ASSERT_EQ(7, decoded.DataLevel()); + ASSERT_EQ(files, decoded.SourceFiles()); +} + +TEST_F(PrimaryKeyIndexSourceMetaTest, FromIndexFileRejectsMissingSourceMeta) { + // Index file without any global index metadata. + IndexFileMeta no_global_index("pk-btree", "index-file-1", /*file_size=*/64, /*row_count=*/100, + /*dv_ranges=*/std::nullopt, /*external_path=*/std::nullopt); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::FromIndexFile(no_global_index)); + + // Global index metadata whose source_meta is null. + std::shared_ptr index_meta = std::make_shared("index-payload", pool_.get()); + GlobalIndexMeta null_source_meta(/*_row_range_start=*/0, /*_row_range_end=*/100, + /*_index_field_id=*/1, /*_extra_field_ids=*/std::nullopt, + index_meta, /*_source_meta=*/nullptr); + IndexFileMeta no_source_meta("pk-btree", "index-file-2", /*file_size=*/64, /*row_count=*/100, + /*dv_ranges=*/std::nullopt, /*external_path=*/std::nullopt, + null_source_meta); + ASSERT_NOK(PrimaryKeyIndexSourceMeta::FromIndexFile(no_source_meta)); +} + +} // namespace paimon::test diff --git a/src/paimon/core/index/pk/primary_key_index_source_policy.h b/src/paimon/core/index/pk/primary_key_index_source_policy.h new file mode 100644 index 000000000..5251bf9a9 --- /dev/null +++ b/src/paimon/core/index/pk/primary_key_index_source_policy.h @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include + +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/manifest/file_source.h" + +namespace paimon { +/// Selects complete compacted data files for source-backed primary-key indexes. +/// +/// Only files produced by compaction on a positive data level are eligible: level 0 files +/// and appended files may still be rewritten or merged, so payloads built over them could +/// not maintain the exact per-level coverage contract. +class PrimaryKeyIndexSourcePolicy { + public: + PrimaryKeyIndexSourcePolicy() = delete; + ~PrimaryKeyIndexSourcePolicy() = delete; + + static bool ShouldWrite(const FileSource& file_source, int32_t level) { + return file_source == FileSource::Compact() && level > 0; + } + + static bool ShouldRead(const DataFileMeta& file) { + if (file.file_source == std::nullopt) { + return false; + } + return ShouldWrite(file.file_source.value(), file.level); + } +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.cpp b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.cpp new file mode 100644 index 000000000..1751eef7f --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.cpp @@ -0,0 +1,112 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h" + +#include +#include +#include +#include +#include + +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_policy.h" + +namespace paimon { +PkSortedBucketIndexState PkSortedBucketIndexState::FromActiveDataFiles( + int32_t field_id, const std::string& index_type, + const std::vector>& active_data_files, + const std::vector>& active_payloads) { + std::map> sources_by_level; + for (const std::shared_ptr& data_file : active_data_files) { + if (data_file != nullptr && PrimaryKeyIndexSourcePolicy::ShouldRead(*data_file)) { + sources_by_level[data_file->level].emplace_back(data_file->file_name, + data_file->row_count); + } + } + for (auto& level_sources : sources_by_level) { + std::sort( + level_sources.second.begin(), level_sources.second.end(), + [](const PrimaryKeyIndexSourceFile& left, const PrimaryKeyIndexSourceFile& right) { + return left.file_name < right.file_name; + }); + } + + // Match payloads against the expected level sources; anything that does not decode or + // does not exactly cover its level is rejected. + std::map>> payloads_by_level; + std::map> payload_metas_by_level; + std::vector> rejected; + for (const std::shared_ptr& payload : active_payloads) { + if (payload == nullptr) { + continue; + } + const std::optional& global_index_meta = payload->GetGlobalIndexMeta(); + if (payload->IndexType() != index_type || global_index_meta == std::nullopt || + global_index_meta->index_field_id != field_id) { + rejected.push_back(payload); + continue; + } + Result source_meta_result = + PrimaryKeyIndexSourceMeta::FromIndexFile(*payload); + if (!source_meta_result.ok()) { + rejected.push_back(payload); + continue; + } + PrimaryKeyIndexSourceMeta source_meta = std::move(source_meta_result).value(); + auto desired = sources_by_level.find(source_meta.DataLevel()); + if (desired == sources_by_level.end() || desired->second != source_meta.SourceFiles()) { + rejected.push_back(payload); + continue; + } + payloads_by_level[source_meta.DataLevel()].push_back(payload); + payload_metas_by_level[source_meta.DataLevel()].push_back(std::move(source_meta)); + } + + std::vector groups; + std::set covered_levels; + for (const auto& level_payloads : payloads_by_level) { + int32_t level = level_payloads.first; + std::optional group; + if (level_payloads.second.size() == 1) { + group = PkSortedIndexGroup::Create(field_id, index_type, sources_by_level[level], + level_payloads.second[0], + payload_metas_by_level[level][0]); + } + if (group != std::nullopt) { + groups.push_back(std::move(group).value()); + covered_levels.insert(level); + } else { + rejected.insert(rejected.end(), level_payloads.second.begin(), + level_payloads.second.end()); + } + } + + std::vector covered; + std::vector uncovered; + for (const auto& level_sources : sources_by_level) { + auto& target = covered_levels.count(level_sources.first) > 0 ? covered : uncovered; + target.insert(target.end(), level_sources.second.begin(), level_sources.second.end()); + } + return PkSortedBucketIndexState(std::move(groups), std::move(covered), std::move(uncovered), + std::move(rejected)); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.h b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.h new file mode 100644 index 000000000..42250288d --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state.h @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pksorted/pk_sorted_index_group.h" +#include "paimon/core/io/data_file_meta.h" + +namespace paimon { +/// Immutable source-backed sorted-index state for one field and bucket. +/// +/// Derives the eligible per-level source sets from the active data files, matches the +/// active payloads against them, and keeps the exact validated groups. Payloads whose +/// source metadata cannot be decoded or does not exactly cover its level are rejected; +/// levels without a valid group stay uncovered and must be scanned normally. +class PkSortedBucketIndexState { + public: + static PkSortedBucketIndexState FromActiveDataFiles( + int32_t field_id, const std::string& index_type, + const std::vector>& active_data_files, + const std::vector>& active_payloads); + + const std::vector& Groups() const { + return groups_; + } + + const std::vector& CoveredSourceFiles() const { + return covered_source_files_; + } + + const std::vector& UncoveredSourceFiles() const { + return uncovered_source_files_; + } + + const std::vector>& RejectedPayloads() const { + return rejected_payloads_; + } + + private: + PkSortedBucketIndexState(std::vector groups, + std::vector covered_source_files, + std::vector uncovered_source_files, + std::vector> rejected_payloads) + : groups_(std::move(groups)), + covered_source_files_(std::move(covered_source_files)), + uncovered_source_files_(std::move(uncovered_source_files)), + rejected_payloads_(std::move(rejected_payloads)) {} + + std::vector groups_; + std::vector covered_source_files_; + std::vector uncovered_source_files_; + std::vector> rejected_payloads_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp new file mode 100644 index 000000000..d85045cac --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_bucket_index_state_test.cpp @@ -0,0 +1,303 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h" + +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" +#include "paimon/core/index/pksorted/pk_sorted_index_group.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/manifest/file_source.h" +#include "paimon/core/stats/simple_stats.h" +#include "paimon/data/timestamp.h" +#include "paimon/memory/bytes.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +class PkSortedBucketIndexStateTest : public ::testing::Test { + public: + std::shared_ptr MakeDataFile(const std::string& file_name, int64_t row_count, + int32_t level, + const std::optional& file_source) const { + return std::make_shared( + file_name, /*file_size=*/1024, row_count, DataFileMeta::EmptyMinKey(), + DataFileMeta::EmptyMaxKey(), SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), + /*min_sequence_number=*/0, /*max_sequence_number=*/1, /*schema_id=*/0, level, + /*extra_files=*/std::vector>(), + /*creation_time=*/Timestamp(0, 0), /*delete_row_count=*/std::nullopt, + /*embedded_index=*/nullptr, file_source, /*value_stats_cols=*/std::nullopt, + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, + /*write_cols=*/std::nullopt); + } + + /// Builds a payload whose source metadata lists the given sources in the given order. + std::shared_ptr MakePayload( + int32_t field_id, const std::string& index_type, int32_t data_level, + const std::vector& sources, int64_t total_row_count, + int64_t row_range_start, int64_t row_range_end) const { + EXPECT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create(data_level, sources)); + EXPECT_OK_AND_ASSIGN(std::shared_ptr source_meta_bytes, + source_meta.Serialize(pool_)); + return MakePayloadWithSourceMetaBytes(field_id, index_type, total_row_count, + row_range_start, row_range_end, source_meta_bytes); + } + + std::shared_ptr MakePayload( + int32_t field_id, const std::string& index_type, int32_t data_level, + const std::vector& sources, int64_t total_row_count) const { + return MakePayload(field_id, index_type, data_level, sources, total_row_count, + /*row_range_start=*/0, /*row_range_end=*/total_row_count - 1); + } + + std::shared_ptr MakePayloadWithSourceMetaBytes( + int32_t field_id, const std::string& index_type, int64_t total_row_count, + int64_t row_range_start, int64_t row_range_end, + const std::shared_ptr& source_meta_bytes) const { + GlobalIndexMeta global_index_meta(row_range_start, row_range_end, field_id, + /*extra_field_ids=*/std::nullopt, + /*index_meta=*/nullptr, source_meta_bytes); + return std::make_shared(index_type, /*file_name=*/"payload.index", + /*file_size=*/2048, total_row_count, + /*dv_ranges=*/std::nullopt, + /*external_path=*/std::nullopt, global_index_meta); + } + + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +TEST_F(PkSortedBucketIndexStateTest, BuildsGroupWhenPayloadMatchesLevelSources) { + // Files are handed over unsorted; the expected source order is sorted by file name. + std::vector> data_files = { + MakeDataFile("b", 200, 5, FileSource::Compact()), + MakeDataFile("a", 100, 5, FileSource::Compact())}; + std::vector expected_sources = {{"a", 100}, {"b", 200}}; + std::shared_ptr payload = + MakePayload(/*field_id=*/7, "btree", /*data_level=*/5, expected_sources, + /*total_row_count=*/300); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_EQ(1, state.Groups().size()); + const PkSortedIndexGroup& group = state.Groups()[0]; + ASSERT_EQ(5, group.DataLevel()); + ASSERT_EQ(300, group.TotalSourceRowCount()); + ASSERT_EQ(expected_sources, group.SourceFiles()); + ASSERT_EQ(payload, group.Payload()); + ASSERT_EQ(expected_sources, state.CoveredSourceFiles()); + ASSERT_TRUE(state.UncoveredSourceFiles().empty()); + ASSERT_TRUE(state.RejectedPayloads().empty()); +} + +TEST_F(PkSortedBucketIndexStateTest, OnlyCompactedFilesAboveLevelZeroAreSources) { + std::vector> data_files = { + MakeDataFile("level0", 10, 0, FileSource::Compact()), + MakeDataFile("appended", 20, 5, FileSource::Append()), + MakeDataFile("unknown_source", 30, 5, std::nullopt), + MakeDataFile("c", 40, 5, FileSource::Compact())}; + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_TRUE(state.CoveredSourceFiles().empty()); + std::vector expected_uncovered = {{"c", 40}}; + ASSERT_EQ(expected_uncovered, state.UncoveredSourceFiles()); + ASSERT_TRUE(state.RejectedPayloads().empty()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithMisorderedSources) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::shared_ptr payload = + MakePayload(7, "btree", 5, {{"b", 200}, {"a", 100}}, 300); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(payload, state.RejectedPayloads()[0]); + ASSERT_TRUE(state.CoveredSourceFiles().empty()); + std::vector expected_uncovered = {{"a", 100}, {"b", 200}}; + ASSERT_EQ(expected_uncovered, state.UncoveredSourceFiles()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithMismatchedSourceRowCount) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::shared_ptr payload = + MakePayload(7, "btree", 5, {{"a", 100}, {"b", 201}}, 301); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(2, state.UncoveredSourceFiles().size()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadsCoveringWrongSourceSet) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::shared_ptr missing_source_payload = + MakePayload(7, "btree", 5, {{"a", 100}}, 100); + std::shared_ptr extra_source_payload = + MakePayload(7, "btree", 5, {{"a", 100}, {"b", 200}, {"c", 50}}, 350); + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + 7, "btree", data_files, {missing_source_payload, extra_source_payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(2, state.RejectedPayloads().size()); + ASSERT_TRUE(state.CoveredSourceFiles().empty()); + ASSERT_EQ(2, state.UncoveredSourceFiles().size()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsBothPayloadsWhenLevelHasTwoCandidates) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::vector sources = {{"a", 100}, {"b", 200}}; + std::shared_ptr first_payload = MakePayload(7, "btree", 5, sources, 300); + std::shared_ptr second_payload = MakePayload(7, "btree", 5, sources, 300); + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + 7, "btree", data_files, {first_payload, second_payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(2, state.RejectedPayloads().size()); + ASSERT_TRUE(state.CoveredSourceFiles().empty()); + ASSERT_EQ(sources, state.UncoveredSourceFiles()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithWrongFieldId) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact())}; + std::shared_ptr payload = + MakePayload(/*field_id=*/8, "btree", 5, {{"a", 100}}, 100); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(1, state.UncoveredSourceFiles().size()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithWrongIndexType) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact())}; + std::shared_ptr payload = MakePayload(7, "bitmap", 5, {{"a", 100}}, 100); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(1, state.UncoveredSourceFiles().size()); +} + +TEST_F(PkSortedBucketIndexStateTest, WrongCandidateDoesNotMaskValidPayload) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact())}; + std::shared_ptr valid_payload = MakePayload(7, "btree", 5, {{"a", 100}}, 100); + std::shared_ptr wrong_payload = + MakePayload(/*field_id=*/8, "btree", 5, {{"a", 100}}, 100); + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + 7, "btree", data_files, {valid_payload, wrong_payload}); + ASSERT_EQ(1, state.Groups().size()); + ASSERT_EQ(valid_payload, state.Groups()[0].Payload()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(wrong_payload, state.RejectedPayloads()[0]); + ASSERT_TRUE(state.UncoveredSourceFiles().empty()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithWrongRowRange) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::vector sources = {{"a", 100}, {"b", 200}}; + // The exclusive end row 300 violates the required inclusive range [0, 299]. + std::shared_ptr payload = MakePayload(7, "btree", 5, sources, 300, + /*row_range_start=*/0, + /*row_range_end=*/300); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(sources, state.UncoveredSourceFiles()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithWrongRowCount) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + std::vector sources = {{"a", 100}, {"b", 200}}; + // The row range is valid but the payload row count 299 differs from the 300 source rows. + std::shared_ptr payload = MakePayload(7, "btree", 5, sources, + /*total_row_count=*/299, + /*row_range_start=*/0, + /*row_range_end=*/299); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(sources, state.UncoveredSourceFiles()); +} + +TEST_F(PkSortedBucketIndexStateTest, RejectsPayloadWithCorruptSourceMeta) { + std::vector> data_files = { + MakeDataFile("a", 100, 5, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact())}; + // Version 1 followed by a truncated data level. + std::shared_ptr corrupt_source_meta = + std::make_shared(std::string("\x00\x00\x00\x01\x00\x00", 6), pool_.get()); + std::shared_ptr payload = + MakePayloadWithSourceMetaBytes(7, "btree", /*total_row_count=*/300, /*row_range_start=*/0, + /*row_range_end=*/299, corrupt_source_meta); + PkSortedBucketIndexState state = + PkSortedBucketIndexState::FromActiveDataFiles(7, "btree", data_files, {payload}); + ASSERT_TRUE(state.Groups().empty()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(payload, state.RejectedPayloads()[0]); + ASSERT_EQ(2, state.UncoveredSourceFiles().size()); +} + +TEST_F(PkSortedBucketIndexStateTest, KeepsValidLevelAndLeavesBrokenLevelUncovered) { + std::vector> data_files = { + MakeDataFile("c", 50, 4, FileSource::Compact()), + MakeDataFile("b", 200, 5, FileSource::Compact()), + MakeDataFile("a", 100, 5, FileSource::Compact())}; + std::shared_ptr valid_payload = MakePayload(7, "btree", 4, {{"c", 50}}, 50); + std::shared_ptr broken_payload = + MakePayload(7, "btree", 5, {{"a", 100}, {"b", 999}}, 1099); + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + 7, "btree", data_files, {valid_payload, broken_payload}); + ASSERT_EQ(1, state.Groups().size()); + ASSERT_EQ(4, state.Groups()[0].DataLevel()); + ASSERT_EQ(valid_payload, state.Groups()[0].Payload()); + std::vector expected_covered = {{"c", 50}}; + ASSERT_EQ(expected_covered, state.CoveredSourceFiles()); + std::vector expected_uncovered = {{"a", 100}, {"b", 200}}; + ASSERT_EQ(expected_uncovered, state.UncoveredSourceFiles()); + ASSERT_EQ(1, state.RejectedPayloads().size()); + ASSERT_EQ(broken_payload, state.RejectedPayloads()[0]); +} + +} // namespace paimon::test diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp b/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp new file mode 100644 index 000000000..f919995b1 --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp @@ -0,0 +1,115 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pksorted/pk_sorted_index_file.h" + +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" +#include "paimon/global_index/global_index_io_meta.h" +#include "paimon/global_index/global_index_writer.h" +#include "paimon/global_index/global_indexer.h" +#include "paimon/global_index/global_indexer_factory.h" + +namespace paimon { +Result> PkSortedIndexFile::Build( + const DataField& field, const std::string& index_type, + const std::map& options, int32_t data_level, + const std::vector& source_files, + const std::shared_ptr& sorted_values, std::vector sorted_ordinals, + const std::shared_ptr& file_writer, bool is_external_path, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create(data_level, source_files)); + int64_t source_row_count = 0; + for (const PrimaryKeyIndexSourceFile& source_file : source_files) { + if (__builtin_add_overflow(source_row_count, source_file.row_count, &source_row_count)) { + return Status::Invalid("Source row count overflows in sorted index build."); + } + } + if (source_row_count <= 0) { + return Status::Invalid("A sorted index group must reference at least one source row."); + } + if (sorted_values == nullptr || sorted_values->length() != source_row_count || + static_cast(sorted_ordinals.size()) != source_row_count) { + return Status::Invalid( + fmt::format("Sorted index input row count {} does not match source row count {}.", + sorted_values == nullptr ? 0 : sorted_values->length(), source_row_count)); + } + std::vector seen_ordinals(static_cast(source_row_count), false); + for (int64_t ordinal : sorted_ordinals) { + if (ordinal < 0 || ordinal >= source_row_count) { + return Status::Invalid( + fmt::format("Row id {} is outside sorted index group row range [0, {}).", ordinal, + source_row_count)); + } + if (seen_ordinals[ordinal]) { + return Status::Invalid(fmt::format("Row id {} appears more than once.", ordinal)); + } + seen_ordinals[ordinal] = true; + } + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, + GlobalIndexerFactory::Get(index_type, options)); + if (indexer == nullptr) { + return Status::Invalid(fmt::format("Index type {} is not registered.", index_type)); + } + auto arrow_field = DataField::ConvertDataFieldToArrowField(field); + auto arrow_schema = arrow::schema({arrow_field}); + ArrowSchema c_arrow_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema, &c_arrow_schema)); + ScopeGuard schema_guard([&]() { ArrowSchemaRelease(&c_arrow_schema); }); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, + indexer->CreateWriter(field.Name(), &c_arrow_schema, file_writer, pool)); + + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr struct_array, + arrow::StructArray::Make({sorted_values}, {field.Name()})); + ::ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, &c_array)); + ScopeGuard array_guard([&]() { ArrowArrayRelease(&c_array); }); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array, std::move(sorted_ordinals))); + PAIMON_ASSIGN_OR_RAISE(std::vector io_metas, writer->Finish()); + if (io_metas.size() != 1) { + return Status::Invalid(fmt::format( + "Sorted index build must produce exactly one payload file, but produced {}.", + io_metas.size())); + } + const GlobalIndexIOMeta& io_meta = io_metas[0]; + + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr source_meta_bytes, source_meta.Serialize(pool)); + std::optional external_path; + if (is_external_path) { + external_path = io_meta.file_path; + } + return std::make_shared( + index_type, PathUtil::GetName(io_meta.file_path), io_meta.file_size, source_row_count, + /*dv_ranges=*/std::nullopt, external_path, + GlobalIndexMeta(0, source_row_count - 1, field.Id(), + /*extra_field_ids=*/std::nullopt, io_meta.metadata, source_meta_bytes)); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_file.h b/src/paimon/core/index/pksorted/pk_sorted_index_file.h new file mode 100644 index 000000000..1c5089e94 --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_file.h @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/global_index/io/global_index_file_writer.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" + +namespace paimon { +/// Builds one source-backed primary-key index payload for an ordered set of physical data +/// files of a single data level. +/// +/// The caller provides all indexed values of the source group as one array sorted by +/// value, together with each value's zero-based ordinal in the ordered source group. The +/// builder writes exactly one payload file and returns its metadata carrying the +/// serialized `PrimaryKeyIndexSourceMeta`, so the payload can later be validated against +/// the active source set of its level. +class PkSortedIndexFile { + public: + PkSortedIndexFile() = delete; + ~PkSortedIndexFile() = delete; + + /// @param field The indexed field. + /// @param index_type The index algorithm identifier, e.g. "btree". + /// @param options The resolved index options (algorithm-prefixed keys included). + /// @param data_level The positive data level covered by the payload. + /// @param source_files The level's active source files ordered by file name. + /// @param sorted_values All indexed values of the source group sorted by value; nulls + /// may appear anywhere. + /// @param sorted_ordinals The group ordinal of each value, aligned with + /// `sorted_values`; every ordinal in `[0, total source rows)` must appear + /// exactly once. + /// @param file_writer The index-directory file writer of the payload's bucket. + /// @param is_external_path Whether `file_writer` resolves to an external index path. + /// @param pool The memory pool used for metadata and index construction. + /// @return Metadata for the single payload file written by the index builder. + static Result> Build( + const DataField& field, const std::string& index_type, + const std::map& options, int32_t data_level, + const std::vector& source_files, + const std::shared_ptr& sorted_values, std::vector sorted_ordinals, + const std::shared_ptr& file_writer, bool is_external_path, + const std::shared_ptr& pool); +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_group.cpp b/src/paimon/core/index/pksorted/pk_sorted_index_group.cpp new file mode 100644 index 000000000..b226dc11d --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_group.cpp @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/index/pksorted/pk_sorted_index_group.h" + +#include +#include + +namespace paimon { +std::optional PkSortedIndexGroup::Create( + int32_t field_id, const std::string& index_type, + const std::vector& expected_sources, + const std::shared_ptr& payload, + const PrimaryKeyIndexSourceMeta& payload_source_meta) { + if (payload == nullptr || expected_sources.empty()) { + return std::nullopt; + } + int64_t source_row_count = 0; + std::set source_names; + for (const PrimaryKeyIndexSourceFile& source_file : expected_sources) { + if (!source_names.insert(source_file.file_name).second) { + return std::nullopt; + } + if (__builtin_add_overflow(source_row_count, source_file.row_count, &source_row_count)) { + return std::nullopt; + } + } + + const std::optional& meta = payload->GetGlobalIndexMeta(); + if (payload_source_meta.SourceFiles() != expected_sources || + index_type != payload->IndexType() || meta == std::nullopt || + meta.value().index_field_id != field_id || meta.value().row_range_start != 0 || + meta.value().row_range_end != source_row_count - 1 || + payload->RowCount() != source_row_count) { + return std::nullopt; + } + return PkSortedIndexGroup(payload_source_meta.DataLevel(), expected_sources, payload, + source_row_count); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_group.h b/src/paimon/core/index/pksorted/pk_sorted_index_group.h new file mode 100644 index 000000000..e202f8956 --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_group.h @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" + +namespace paimon { +/// The single validated payload group which indexes one complete data level. +/// +/// A group is only created when the payload provably covers the current active source set +/// of its data level: exactly one payload, unique source names, source order / file names / +/// row counts identical to the expected level sources, matching index type and field id, +/// a row range of exactly `[0, total source rows - 1]` and a payload row count equal to the +/// source row count sum. Anything else must be treated as uncovered. +class PkSortedIndexGroup { + public: + /// Validates one payload against the expected level sources; returns `std::nullopt` + /// when any coverage condition fails. + static std::optional Create( + int32_t field_id, const std::string& index_type, + const std::vector& expected_sources, + const std::shared_ptr& payload, + const PrimaryKeyIndexSourceMeta& payload_source_meta); + + int32_t DataLevel() const { + return data_level_; + } + + const std::vector& SourceFiles() const { + return source_files_; + } + + const std::shared_ptr& Payload() const { + return payload_; + } + + int64_t TotalSourceRowCount() const { + return total_source_row_count_; + } + + private: + PkSortedIndexGroup(int32_t data_level, std::vector source_files, + std::shared_ptr payload, int64_t total_source_row_count) + : data_level_(data_level), + source_files_(std::move(source_files)), + payload_(std::move(payload)), + total_source_row_count_(total_source_row_count) {} + + int32_t data_level_; + std::vector source_files_; + std::shared_ptr payload_; + int64_t total_source_row_count_; +}; + +} // namespace paimon diff --git a/src/paimon/core/operation/raw_file_split_read.cpp b/src/paimon/core/operation/raw_file_split_read.cpp index eabe84268..8bb59987e 100644 --- a/src/paimon/core/operation/raw_file_split_read.cpp +++ b/src/paimon/core/operation/raw_file_split_read.cpp @@ -18,12 +18,14 @@ #include "paimon/core/operation/raw_file_split_read.h" +#include #include #include #include #include "arrow/c/abi.h" #include "arrow/c/bridge.h" +#include "fmt/format.h" #include "paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/concat_batch_reader.h" @@ -32,6 +34,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" #include "paimon/core/deletionvectors/deletion_vector.h" +#include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/file_index_evaluator.h" #include "paimon/core/operation/internal_read_context.h" @@ -64,6 +67,21 @@ RawFileSplitRead::RawFileSplitRead(const std::shared_ptr& Result> RawFileSplitRead::CreateReader( const std::shared_ptr& split) { + if (auto indexed_split = std::dynamic_pointer_cast(split)) { + PAIMON_RETURN_NOT_OK(indexed_split->Validate()); + const std::shared_ptr& inner_split = indexed_split->GetDataSplit(); + auto inner_split_impl = std::dynamic_pointer_cast(inner_split); + if (!inner_split_impl) { + return Status::Invalid("cannot cast indexed inner split to data_split"); + } + if (inner_split_impl->DataFiles().size() != 1) { + return Status::Invalid( + "indexed splits with file-local row ranges must contain exactly one file"); + } + return CreateReader(inner_split_impl->Partition(), inner_split_impl->Bucket(), + inner_split_impl->DataFiles(), inner_split_impl->DeletionFiles(), + indexed_split->RowRanges()); + } auto data_split = std::dynamic_pointer_cast(split); if (!data_split) { return Status::Invalid("cannot cast split to data_split in RawFileSplitRead"); @@ -75,7 +93,7 @@ Result> RawFileSplitRead::CreateReader( Result> RawFileSplitRead::CreateReader( const BinaryRow& partition, int32_t bucket, const std::vector>& data_files, - DeletionVector::Factory dv_factory) { + DeletionVector::Factory dv_factory, const std::optional>& local_row_ranges) { const auto& predicate = context_->GetPredicate(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, path_factory_->CreateDataFilePathFactory(partition, bucket)); @@ -83,7 +101,7 @@ Result> RawFileSplitRead::CreateReader( PAIMON_ASSIGN_OR_RAISE( std::vector> raw_file_readers, CreateRawFileReaders(partition, data_files, raw_read_schema_, predicate, dv_factory, - /*row_ranges=*/{}, data_file_path_factory, + local_row_ranges, data_file_path_factory, /*extra_format_options=*/{})); auto raw_readers = @@ -98,10 +116,19 @@ Result> RawFileSplitRead::CreateReader( const BinaryRow& partition, int32_t bucket, const std::vector>& data_files, const std::vector>& deletion_files) { + return CreateReader(partition, bucket, data_files, deletion_files, + /*local_row_ranges=*/std::nullopt); +} + +Result> RawFileSplitRead::CreateReader( + const BinaryRow& partition, int32_t bucket, + const std::vector>& data_files, + const std::vector>& deletion_files, + const std::optional>& local_row_ranges) { auto dv_factory = DeletionVector::CreateFactory( options_.GetFileSystem(), DeletionVector::CreateDeletionFileMap(data_files, deletion_files), pool_); - return CreateReader(partition, bucket, data_files, dv_factory); + return CreateReader(partition, bucket, data_files, dv_factory, local_row_ranges); } Result RawFileSplitRead::Match(const std::shared_ptr& split, @@ -152,6 +179,28 @@ Result> RawFileSplitRead::ApplyIndexAndDvReader PAIMON_ASSIGN_OR_RAISE(selection, bitmap_file_index->GetBitmap()); } + // narrow the selection to the file-local row positions of an indexed split + std::optional ranges_selection; + if (ranges != std::nullopt) { + RoaringBitmap32 ranges_bitmap; + for (const Range& range : ranges.value()) { + if (range.from < 0 || range.to < range.from || + range.to >= std::numeric_limits::max()) { + return Status::Invalid( + fmt::format("Invalid file-local row range [{}, {}] for file {}.", range.from, + range.to, file->file_name)); + } + ranges_bitmap.AddRange(static_cast(range.from), + static_cast(range.to + 1)); + } + if (selection != nullptr) { + ranges_selection = RoaringBitmap32::And(*selection, ranges_bitmap); + } else { + ranges_selection = std::move(ranges_bitmap); + } + selection = &ranges_selection.value(); + } + // prepare deletion bitmap for deletion vector std::shared_ptr deletion_vector; if (dv_factory) { diff --git a/src/paimon/core/operation/raw_file_split_read.h b/src/paimon/core/operation/raw_file_split_read.h index 1580cd848..9b87dc835 100644 --- a/src/paimon/core/operation/raw_file_split_read.h +++ b/src/paimon/core/operation/raw_file_split_read.h @@ -19,8 +19,7 @@ #pragma once #include -#include -#include +#include #include #include "paimon/core/core_options.h" @@ -65,16 +64,26 @@ class RawFileSplitRead : public AbstractSplitRead { const std::shared_ptr& memory_pool, const std::shared_ptr& executor); + /// Also accepts an `IndexedSplit` over a single-file data split, in which case its row + /// ranges narrow the read to the given file-local physical positions. Result> CreateReader(const std::shared_ptr& split) override; Result> CreateReader( const BinaryRow& partition, int32_t bucket, const std::vector>& files, const std::vector>& deletion_files); + /// Reads with an optional selection of file-local row positions. The ranges apply to + /// every file of the split, so callers pass them only for single-file splits. Result> CreateReader( const BinaryRow& partition, int32_t bucket, const std::vector>& files, - DeletionVector::Factory dv_factory); + const std::vector>& deletion_files, + const std::optional>& local_row_ranges); + + Result> CreateReader( + const BinaryRow& partition, int32_t bucket, + const std::vector>& files, DeletionVector::Factory dv_factory, + const std::optional>& local_row_ranges = std::nullopt); Result Match(const std::shared_ptr& split, bool force_keep_delete) const override; diff --git a/src/paimon/core/table/source/fallback_data_split_test.cpp b/src/paimon/core/table/source/fallback_data_split_test.cpp index 18fed0cf7..80d290560 100644 --- a/src/paimon/core/table/source/fallback_data_split_test.cpp +++ b/src/paimon/core/table/source/fallback_data_split_test.cpp @@ -27,10 +27,12 @@ #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" +#include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/source/data_split_impl.h" +#include "paimon/core/table/source/fallback_table_read.h" #include "paimon/data/timestamp.h" #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" @@ -41,6 +43,41 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { +class TrackingTableRead : public TableRead { + public: + explicit TrackingTableRead(const std::shared_ptr& pool) : TableRead(pool) {} + + Result> CreateReader( + const std::shared_ptr& split) override { + last_split_ = split; + return std::unique_ptr(); + } + + std::shared_ptr last_split_; +}; +} // namespace + +TEST(FallbackTableReadTest, RoutesIndexedSplitToMainTable) { + std::shared_ptr pool = GetDefaultPool(); + auto main_table = std::make_unique(pool); + auto fallback_table = std::make_unique(pool); + TrackingTableRead* main_table_ptr = main_table.get(); + TrackingTableRead* fallback_table_ptr = fallback_table.get(); + FallbackTableRead table_read(std::move(main_table), std::move(fallback_table), pool); + + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, /*bucket_path=*/"", + /*data_files=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr data_split, + builder.IsStreaming(false).RawConvertible(true).Build()); + std::shared_ptr indexed_split = + std::make_shared(data_split, std::vector()); + + ASSERT_OK(table_read.CreateReader(indexed_split)); + ASSERT_EQ(indexed_split, main_table_ptr->last_split_); + ASSERT_EQ(nullptr, fallback_table_ptr->last_split_); +} + TEST(FallbackDataSplitTest, TestDeserialize) { std::string file_name = paimon::test::GetDataDir() + "/parquet/append_table_with_append_pt_branch.db/" diff --git a/src/paimon/core/table/source/fallback_table_read.cpp b/src/paimon/core/table/source/fallback_table_read.cpp index 6ca44215e..5f5957609 100644 --- a/src/paimon/core/table/source/fallback_table_read.cpp +++ b/src/paimon/core/table/source/fallback_table_read.cpp @@ -21,6 +21,7 @@ #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/fallback_data_split.h" +#include "paimon/global_index/indexed_split.h" #include "paimon/status.h" #include "paimon/table/source/data_split.h" @@ -35,6 +36,9 @@ Result> FallbackTableRead::CreateReader( return main_table_->CreateReader(fallback_data_split->GetSplit()); } } + if (std::dynamic_pointer_cast(split) != nullptr) { + return main_table_->CreateReader(split); + } auto data_split = std::dynamic_pointer_cast(split); if (!data_split) { return Status::Invalid("cannot cast split to data split"); diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 23b902260..ad085fb3f 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -21,6 +21,7 @@ #include +#include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" #include "paimon/core/table/source/data_split_impl.h" @@ -74,7 +75,25 @@ void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { - auto data_split = std::dynamic_pointer_cast(split); + std::shared_ptr dispatch_split = split; + if (auto indexed_split = std::dynamic_pointer_cast(split)) { + // A primary-key sorted-index split narrows one raw-readable file to file-local row + // positions. If the raw read cannot serve the inner split, fall back to reading the + // whole file: the index only narrows the scan, so the unnarrowed read stays correct. + const std::shared_ptr& inner_split = indexed_split->GetDataSplit(); + for (const auto& read : split_reads_) { + auto* raw_read = dynamic_cast(read.get()); + if (raw_read == nullptr) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(bool matched, read->Match(inner_split, force_keep_delete_)); + if (matched) { + return read->CreateReader(indexed_split); + } + } + dispatch_split = inner_split; + } + auto data_split = std::dynamic_pointer_cast(dispatch_split); if (!data_split) { return Status::Invalid("split cannot be casted to DataSplit"); } diff --git a/src/paimon/core/table/source/primary_key_index_batch_scan.cpp b/src/paimon/core/table/source/primary_key_index_batch_scan.cpp new file mode 100644 index 000000000..9425c1a06 --- /dev/null +++ b/src/paimon/core/table/source/primary_key_index_batch_scan.cpp @@ -0,0 +1,298 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/table/source/primary_key_index_batch_scan.h" + +#include +#include +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/core/index/index_file_handler.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/table/source/plan_impl.h" +#include "paimon/core/table/source/primary_key_sorted_index_result.h" +#include "paimon/core/table/source/primary_key_sorted_index_scan.h" +#include "paimon/core/table/source/snapshot/snapshot_reader.h" +#include "paimon/core/utils/index_file_path_factories.h" +#include "paimon/core/utils/snapshot_manager.h" +#include "paimon/executor.h" +#include "paimon/predicate/compound_predicate.h" +#include "paimon/predicate/leaf_predicate.h" +#include "paimon/predicate/predicate_builder.h" + +namespace paimon { +namespace { +Result> CreateGlobalIndexExecutor(const CoreOptions& core_options) { + uint32_t thread_num = std::thread::hardware_concurrency(); + std::optional configured_thread_num = core_options.GetGlobalIndexThreadNum(); + if (configured_thread_num) { + if (configured_thread_num.value() <= 0) { + return Status::Invalid(fmt::format("invalid global index thread number {}", + configured_thread_num.value())); + } + thread_num = static_cast(configured_thread_num.value()); + } else if (thread_num == 0) { + thread_num = 1; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr executor, CreateDefaultExecutor(thread_num)); + return executor; +} + +/// Restricts a predicate to leaves over the indexed fields: an AND keeps its convertible +/// children, an OR is only kept when every child is convertible, and everything else is +/// dropped. A null return means no part of the predicate can use the index. +Result> ProjectToIndexedFields( + const std::shared_ptr& predicate, const std::set& indexed_fields) { + if (predicate == nullptr) { + return std::shared_ptr(nullptr); + } + if (auto leaf_predicate = std::dynamic_pointer_cast(predicate)) { + if (indexed_fields.count(leaf_predicate->FieldName()) > 0) { + return predicate; + } + return std::shared_ptr(nullptr); + } + auto compound_predicate = std::dynamic_pointer_cast(predicate); + if (compound_predicate == nullptr) { + return std::shared_ptr(nullptr); + } + bool is_and = compound_predicate->GetFunction().GetType() == Function::Type::AND; + bool is_or = compound_predicate->GetFunction().GetType() == Function::Type::OR; + if (!is_and && !is_or) { + return std::shared_ptr(nullptr); + } + std::vector> converted_children; + for (const std::shared_ptr& child : compound_predicate->Children()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr converted_child, + ProjectToIndexedFields(child, indexed_fields)); + if (converted_child != nullptr) { + converted_children.push_back(std::move(converted_child)); + } else if (is_or) { + return std::shared_ptr(nullptr); + } + } + if (converted_children.empty()) { + return std::shared_ptr(nullptr); + } + if (converted_children.size() == 1) { + return converted_children[0]; + } + if (is_and) { + return PredicateBuilder::And(converted_children); + } + return PredicateBuilder::Or(converted_children); +} + +void FlattenChildren(const std::shared_ptr& compound_predicate, + std::vector>* flattened) { + for (const std::shared_ptr& child : compound_predicate->Children()) { + auto compound_child = std::dynamic_pointer_cast(child); + if (compound_child != nullptr && compound_child->GetFunction().GetType() == + compound_predicate->GetFunction().GetType()) { + FlattenChildren(compound_child, flattened); + } else { + flattened->push_back(child); + } + } +} + +/// A predicate is null-rejecting when it cannot match a row whose tested field is null. +/// Under SQL three-valued logic every comparison and match predicate rejects null; only +/// IS NULL accepts it, and IS NOT NULL is the predicate being pruned. +bool IsNullRejecting(const std::shared_ptr& predicate) { + auto leaf_predicate = std::dynamic_pointer_cast(predicate); + if (leaf_predicate == nullptr) { + return false; + } + switch (leaf_predicate->GetFunction().GetType()) { + case Function::Type::EQUAL: + case Function::Type::NOT_EQUAL: + case Function::Type::GREATER_THAN: + case Function::Type::GREATER_OR_EQUAL: + case Function::Type::LESS_THAN: + case Function::Type::LESS_OR_EQUAL: + case Function::Type::IN: + case Function::Type::NOT_IN: + case Function::Type::STARTS_WITH: + case Function::Type::ENDS_WITH: + case Function::Type::CONTAINS: + case Function::Type::LIKE: + return true; + default: + return false; + } +} + +bool IsIsNotNull(const std::shared_ptr& predicate) { + auto leaf_predicate = std::dynamic_pointer_cast(predicate); + return leaf_predicate != nullptr && + leaf_predicate->GetFunction().GetType() == Function::Type::IS_NOT_NULL; +} + +/// Flattens nested same-function compounds and, inside an AND, removes `f IS NOT NULL` +/// leaves made redundant by a null-rejecting sibling on the same field. Pruning must not +/// consider `f IS NULL` as constraining: dropping IS NOT NULL from +/// "f IS NULL AND f IS NOT NULL" would turn the empty result into the set of null rows. +Result> NormalizePredicate(const std::shared_ptr& predicate) { + auto compound_predicate = std::dynamic_pointer_cast(predicate); + if (compound_predicate == nullptr) { + return predicate; + } + std::vector> children; + FlattenChildren(compound_predicate, &children); + + bool is_and = compound_predicate->GetFunction().GetType() == Function::Type::AND; + if (is_and) { + std::set constrained_fields; + for (const std::shared_ptr& child : children) { + if (IsNullRejecting(child)) { + constrained_fields.insert( + std::dynamic_pointer_cast(child)->FieldName()); + } + } + if (!constrained_fields.empty()) { + std::vector> pruned; + pruned.reserve(children.size()); + for (const std::shared_ptr& child : children) { + if (IsIsNotNull(child) && + constrained_fields.count( + std::dynamic_pointer_cast(child)->FieldName()) > 0) { + continue; + } + pruned.push_back(child); + } + children = std::move(pruned); + } + } + + std::vector> normalized_children; + normalized_children.reserve(children.size()); + for (const std::shared_ptr& child : children) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_child, + NormalizePredicate(child)); + normalized_children.push_back(std::move(normalized_child)); + } + if (normalized_children.size() == 1) { + return normalized_children[0]; + } + if (is_and) { + return PredicateBuilder::And(normalized_children); + } + return PredicateBuilder::Or(normalized_children); +} +} // namespace + +Result> PrimaryKeyIndexBatchScan::Create( + const std::shared_ptr& snapshot_reader, + std::unique_ptr&& batch_scan, + const std::shared_ptr& table_schema, + const std::shared_ptr& path_factory, const CoreOptions& core_options, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*table_schema)); + return std::unique_ptr(new PrimaryKeyIndexBatchScan( + snapshot_reader, std::move(batch_scan), table_schema, path_factory, core_options, pool, + definitions.ScalarDefinitions())); +} + +Result> PrimaryKeyIndexBatchScan::CreatePlan() { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_plan, batch_scan_->CreatePlan()); + if (!core_options_.GlobalIndexEnabled() || scalar_definitions_.empty() || + data_plan->SnapshotId() == std::nullopt || data_plan->Splits().empty()) { + return data_plan; + } + + std::set indexed_fields; + std::set indexed_field_ids; + for (const PrimaryKeyIndexDefinition& definition : scalar_definitions_) { + indexed_fields.insert(definition.Column()); + indexed_field_ids.insert(definition.FieldId()); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr index_predicate, + ProjectToIndexedFields(batch_scan_->GetNonPartitionPredicate(), indexed_fields)); + if (index_predicate == nullptr) { + return data_plan; + } + PAIMON_ASSIGN_OR_RAISE(index_predicate, NormalizePredicate(index_predicate)); + + std::vector> data_splits; + data_splits.reserve(data_plan->Splits().size()); + for (const std::shared_ptr& split : data_plan->Splits()) { + auto data_split = std::dynamic_pointer_cast(split); + if (data_split == nullptr || data_split->IsStreaming()) { + return data_plan; + } + data_splits.push_back(std::move(data_split)); + } + + int64_t snapshot_id = data_plan->SnapshotId().value(); + const std::shared_ptr& snapshot_manager = + snapshot_reader_->GetSnapshotManager(); + Result snapshot_result = snapshot_manager->LoadSnapshot(snapshot_id); + if (!snapshot_result.ok()) { + return data_plan; + } + + const std::unique_ptr& index_file_handler = + snapshot_reader_->GetIndexFileHandler(); + if (index_file_handler == nullptr) { + return data_plan; + } + std::function(const IndexManifestEntry&)> entry_filter = + [&indexed_field_ids](const IndexManifestEntry& entry) -> Result { + if (!(entry.kind == FileKind::Add()) || entry.index_file == nullptr) { + return false; + } + const std::optional& meta = entry.index_file->GetGlobalIndexMeta(); + return meta != std::nullopt && meta.value().source_meta != nullptr && + indexed_field_ids.count(meta.value().index_field_id) > 0; + }; + PAIMON_ASSIGN_OR_RAISE(std::vector index_entries, + index_file_handler->Scan(snapshot_result.value(), entry_filter)); + + PAIMON_ASSIGN_OR_RAISE(PrimaryKeySortedIndexScan::Plan index_plan, + PrimaryKeySortedIndexScan::CreatePlan( + snapshot_id, data_splits, scalar_definitions_, index_entries)); + bool has_index_group = std::any_of( + index_plan.Files().begin(), index_plan.Files().end(), + [](const PrimaryKeySortedIndexScan::FilePlan& file) { return !file.Groups().empty(); }); + if (!has_index_group) { + return data_plan; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr executor, + CreateGlobalIndexExecutor(core_options_)); + PrimaryKeySortedIndexScan::ReaderFactory reader_factory = + PrimaryKeySortedIndexScan::MakeReaderFactory( + core_options_.GetFileSystem(), std::make_shared(path_factory_), + table_schema_, pool_, executor); + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeySortedIndexScan::EvaluatedPlan evaluated_plan, + PrimaryKeySortedIndexScan::Evaluate(index_plan, table_schema_, index_predicate, + scalar_definitions_, reader_factory)); + PAIMON_ASSIGN_OR_RAISE(std::vector> splits, + PrimaryKeySortedIndexResult::ToSplits(evaluated_plan)); + return std::make_shared(data_plan->SnapshotId(), splits); +} + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_index_batch_scan.h b/src/paimon/core/table/source/primary_key_index_batch_scan.h new file mode 100644 index 000000000..bf5284246 --- /dev/null +++ b/src/paimon/core/table/source/primary_key_index_batch_scan.h @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/table/source/abstract_table_scan.h" +#include "paimon/core/table/source/data_table_batch_scan.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/result.h" + +namespace paimon { +/// Batch scan for primary-key tables with source-backed scalar index definitions. +/// +/// Wraps the ordinary batch scan: the data plan is computed first, then the part of the +/// scan predicate that touches indexed fields is evaluated against the validated payload +/// groups of the plan's snapshot, and covered files are narrowed to indexed splits with +/// file-local row ranges. Files without trustworthy coverage keep their normal scan; the +/// reader still applies deletion vectors and the complete original predicate. +class PrimaryKeyIndexBatchScan : public AbstractTableScan { + public: + static Result> Create( + const std::shared_ptr& snapshot_reader, + std::unique_ptr&& batch_scan, + const std::shared_ptr& table_schema, + const std::shared_ptr& path_factory, const CoreOptions& core_options, + const std::shared_ptr& pool); + + Result> CreatePlan() override; + + private: + PrimaryKeyIndexBatchScan(const std::shared_ptr& snapshot_reader, + std::unique_ptr&& batch_scan, + const std::shared_ptr& table_schema, + const std::shared_ptr& path_factory, + const CoreOptions& core_options, + const std::shared_ptr& pool, + std::vector scalar_definitions) + : AbstractTableScan(core_options, snapshot_reader), + batch_scan_(std::move(batch_scan)), + table_schema_(table_schema), + path_factory_(path_factory), + pool_(pool), + scalar_definitions_(std::move(scalar_definitions)) {} + + std::unique_ptr batch_scan_; + std::shared_ptr table_schema_; + std::shared_ptr path_factory_; + std::shared_ptr pool_; + std::vector scalar_definitions_; +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_sorted_index_result.cpp b/src/paimon/core/table/source/primary_key_sorted_index_result.cpp new file mode 100644 index 000000000..2814b3325 --- /dev/null +++ b/src/paimon/core/table/source/primary_key_sorted_index_result.cpp @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/table/source/primary_key_sorted_index_result.h" + +#include +#include +#include +#include + +#include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/table/source/deletion_file.h" + +namespace paimon { +namespace { +struct RangeConversion { + bool use_index; + std::vector ranges; +}; + +Result> ToSingleFileSplit( + const PrimaryKeySortedIndexScan::FilePlan& file) { + const std::shared_ptr& source = file.SourceSplit(); + std::vector> data_files{file.DataFile()}; + DataSplitImpl::Builder builder(source->Partition(), source->Bucket(), source->BucketPath(), + std::move(data_files)); + builder.WithSnapshot(source->SnapshotId()) + .WithTotalBuckets(source->TotalBuckets()) + .IsStreaming(false) + .RawConvertible(source->RawConvertible()); + if (!source->DeletionFiles().empty()) { + builder.WithDataDeletionFiles({source->DeletionFiles()[file.FileIndex()]}); + } + return builder.Build(); +} + +/// Converts sorted file-local positions to merged ranges. Sets `use_index` to false when a +/// position is invalid or the result is over-fragmented, in which case the file must fall back +/// to a normal scan. +Result ToRanges(const GlobalIndexResult& result, int64_t row_count) { + std::vector ranges; + int64_t from = -1; + int64_t to = -1; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + result.CreateIterator()); + while (iterator->HasNext()) { + int64_t position = iterator->Next(); + if (position < 0 || position >= row_count || + position >= std::numeric_limits::max()) { + return RangeConversion{/*use_index=*/false, {}}; + } + if (from < 0) { + from = position; + } else if (position != to + 1) { + if (ranges.size() >= + static_cast(PrimaryKeySortedIndexResult::kMaxIndexedRangesPerFile)) { + return RangeConversion{/*use_index=*/false, {}}; + } + ranges.emplace_back(from, to); + from = position; + } + to = position; + } + if (ranges.size() >= + static_cast(PrimaryKeySortedIndexResult::kMaxIndexedRangesPerFile)) { + return RangeConversion{/*use_index=*/false, {}}; + } + ranges.emplace_back(from, to); + return RangeConversion{/*use_index=*/true, std::move(ranges)}; +} +} // namespace + +Result>> PrimaryKeySortedIndexResult::ToSplits( + const PrimaryKeySortedIndexScan::EvaluatedPlan& evaluated_plan) { + std::map preserve_raw_splits; + for (const PrimaryKeySortedIndexScan::EvaluatedFile& evaluated_file : evaluated_plan.Files()) { + const std::shared_ptr& source_split = evaluated_file.File().SourceSplit(); + if (!source_split->RawConvertible()) { + continue; + } + auto iter = preserve_raw_splits.emplace(source_split.get(), true).first; + if (evaluated_file.IndexResult() != nullptr) { + iter->second = false; + } + } + + std::vector> splits; + std::set preserved_splits; + for (const PrimaryKeySortedIndexScan::EvaluatedFile& evaluated_file : evaluated_plan.Files()) { + const PrimaryKeySortedIndexScan::FilePlan& file = evaluated_file.File(); + const std::shared_ptr& source_split = file.SourceSplit(); + if (!source_split->RawConvertible() || preserve_raw_splits[source_split.get()]) { + // Preserve the planner's bin packing when the split cannot be read file by file + // or no file in the split has a usable index result. + if (preserved_splits.insert(source_split.get()).second) { + splits.push_back(source_split); + } + continue; + } + + const std::shared_ptr& result = evaluated_file.IndexResult(); + if (result == nullptr) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr fallback_split, + ToSingleFileSplit(file)); + splits.push_back(std::move(fallback_split)); + continue; + } + + PAIMON_ASSIGN_OR_RAISE(bool is_empty, result->IsEmpty()); + if (is_empty) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(RangeConversion range_conversion, + ToRanges(*result, file.DataFile()->row_count)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr single_file_split, + ToSingleFileSplit(file)); + if (!range_conversion.use_index) { + // The index returned an invalid or over-fragmented row position set; fall back + // to a normal scan for this file. + splits.push_back(std::move(single_file_split)); + } else { + splits.push_back(std::make_shared(std::move(single_file_split), + std::move(range_conversion.ranges), + std::vector())); + } + } + return splits; +} + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_sorted_index_result.h b/src/paimon/core/table/source/primary_key_sorted_index_result.h new file mode 100644 index 000000000..8de494ad8 --- /dev/null +++ b/src/paimon/core/table/source/primary_key_sorted_index_result.h @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/core/table/source/primary_key_sorted_index_scan.h" +#include "paimon/table/source/split.h" + +namespace paimon { +/// Converts an evaluated primary-key sorted-index plan into scan splits addressed by +/// physical data-file row positions. +/// +/// Files whose index result is missing or untrustworthy keep a normal single-file scan, +/// files with an empty result are omitted, and files with a valid result become indexed +/// splits carrying their file-local row ranges next to the aligned deletion file. Source +/// splits that are not raw-convertible are preserved unchanged. +class PrimaryKeySortedIndexResult { + public: + /// The fragmentation guard of the Java implementation: a file whose index result needs + /// more ranges falls back to a normal scan. + static constexpr int32_t kMaxIndexedRangesPerFile = 4096; + + PrimaryKeySortedIndexResult() = delete; + ~PrimaryKeySortedIndexResult() = delete; + + static Result>> ToSplits( + const PrimaryKeySortedIndexScan::EvaluatedPlan& evaluated_plan); +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_sorted_index_scan.cpp b/src/paimon/core/table/source/primary_key_sorted_index_scan.cpp new file mode 100644 index 000000000..8b56625d4 --- /dev/null +++ b/src/paimon/core/table/source/primary_key_sorted_index_scan.cpp @@ -0,0 +1,568 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/table/source/primary_key_sorted_index_scan.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/global_index/global_index_evaluator_impl.h" +#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h" +#include "paimon/core/manifest/file_kind.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/global_index/global_index_io_meta.h" +#include "paimon/global_index/global_indexer.h" +#include "paimon/global_index/global_indexer_factory.h" +#include "paimon/global_index/io/global_index_file_reader.h" +#include "paimon/predicate/predicate_utils.h" + +namespace paimon { +namespace { +using BucketKey = std::pair; + +enum class QueryOperation { + IS_NOT_NULL, + IS_NULL, + EQUAL, + NOT_EQUAL, + LESS_THAN, + LESS_OR_EQUAL, + GREATER_THAN, + GREATER_OR_EQUAL, + IN, + NOT_IN, + STARTS_WITH, + ENDS_WITH, + CONTAINS, + LIKE, +}; + +struct QueryKey { + QueryOperation operation; + std::vector literals; + + bool operator==(const QueryKey& other) const { + return operation == other.operation && literals == other.literals; + } +}; + +/// Shares one group payload reader and its group-scope query results across all source +/// files of the group; localizes group ordinals to file-local physical positions using the +/// ordered source row-count prefix. +class SharedGroupReader { + public: + using UnderlyingReaderFactory = std::function>()>; + + SharedGroupReader(const std::shared_ptr& group, + UnderlyingReaderFactory reader_factory) + : group_(group), reader_factory_(std::move(reader_factory)) { + const std::vector& source_files = group->SourceFiles(); + source_offsets_.reserve(source_files.size() + 1); + source_offsets_.push_back(0); + for (const PrimaryKeyIndexSourceFile& source_file : source_files) { + source_offsets_.push_back(source_offsets_.back() + source_file.row_count); + } + } + + const std::shared_ptr& Group() const { + return group_; + } + + /// Runs one group-scope query with caching; equal queries evaluate exactly once. + Result> Query( + const QueryKey& key, + const std::function>(GlobalIndexReader*)>& + query) { + for (const auto& cached : query_cache_) { + if (cached.first == key) { + if (!cached.second.status.ok()) { + return cached.second.status; + } + return cached.second.result; + } + } + Result> result = RunQuery(query); + CachedQuery cached_query; + if (result.ok()) { + cached_query.result = result.value(); + } else { + cached_query.status = result.status(); + } + query_cache_.emplace_back(key, cached_query); + return result; + } + + /// Restricts one group-scope result to the local row positions of `source_index`. + /// Any out-of-range group ordinal fails the localization so that every covered file + /// of this group falls back to a normal scan; a poison marker would not survive the + /// AND/OR combination of results from other indexes. + Result> Localize( + const std::shared_ptr& result, size_t source_index) { + if (result == nullptr) { + return std::shared_ptr(nullptr); + } + assert(source_index + 1 < source_offsets_.size()); + auto localized = localized_cache_.find(result.get()); + if (localized == localized_cache_.end()) { + PAIMON_ASSIGN_OR_RAISE(std::vector> partitions, + PartitionBySource(result)); + localized = localized_cache_.emplace(result.get(), std::move(partitions)).first; + } + return localized->second[source_index]; + } + + private: + struct CachedQuery { + Status status; + std::shared_ptr result; + }; + + Result> RunQuery( + const std::function>(GlobalIndexReader*)>& + query) { + if (!reader_status_.ok()) { + return reader_status_; + } + if (reader_ == nullptr) { + Result> reader_result = reader_factory_(); + if (!reader_result.ok()) { + reader_status_ = reader_result.status(); + return reader_status_; + } + reader_ = reader_result.value(); + if (reader_ == nullptr) { + // The index type has no usable reader; keep normal scan semantics. + return std::shared_ptr(nullptr); + } + } + return query(reader_.get()); + } + + Result>> PartitionBySource( + const std::shared_ptr& result) { + size_t source_count = source_offsets_.size() - 1; + std::vector partitions(source_count); + int64_t total_row_count = source_offsets_.back(); + size_t source_index = 0; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + result->CreateIterator()); + while (iterator->HasNext()) { + int64_t position = iterator->Next(); + if (position < 0 || position >= total_row_count) { + return Status::Invalid(fmt::format( + "Sorted index returned group ordinal {} outside the source row range " + "[0, {}).", + position, total_row_count)); + } + while (position >= source_offsets_[source_index + 1]) { + source_index++; + } + partitions[source_index].Add(position - source_offsets_[source_index]); + } + std::vector> localized; + localized.reserve(source_count); + for (RoaringBitmap64& partition : partitions) { + auto bitmap = std::make_shared(std::move(partition)); + localized.push_back(std::make_shared( + [bitmap]() -> Result { return *bitmap; })); + } + return localized; + } + + std::shared_ptr group_; + UnderlyingReaderFactory reader_factory_; + std::vector source_offsets_; + std::vector> query_cache_; + std::unordered_map>> + localized_cache_; + std::shared_ptr reader_; + Status reader_status_; +}; + +/// Restricts merged source-group ordinals to one source file's local row positions. +class FileLocalGroupReader : public GlobalIndexReader { + public: + FileLocalGroupReader(std::shared_ptr shared_reader, size_t source_index) + : shared_reader_(std::move(shared_reader)), source_index_(source_index) {} + + Result> VisitIsNotNull() override { + return Query({QueryOperation::IS_NOT_NULL, {}}, + [](GlobalIndexReader* reader) { return reader->VisitIsNotNull(); }); + } + + Result> VisitIsNull() override { + return Query({QueryOperation::IS_NULL, {}}, + [](GlobalIndexReader* reader) { return reader->VisitIsNull(); }); + } + + Result> VisitEqual(const Literal& literal) override { + return Query({QueryOperation::EQUAL, {literal}}, + [&literal](GlobalIndexReader* reader) { return reader->VisitEqual(literal); }); + } + + Result> VisitNotEqual(const Literal& literal) override { + return Query({QueryOperation::NOT_EQUAL, {literal}}, [&literal](GlobalIndexReader* reader) { + return reader->VisitNotEqual(literal); + }); + } + + Result> VisitLessThan(const Literal& literal) override { + return Query({QueryOperation::LESS_THAN, {literal}}, [&literal](GlobalIndexReader* reader) { + return reader->VisitLessThan(literal); + }); + } + + Result> VisitLessOrEqual(const Literal& literal) override { + return Query( + {QueryOperation::LESS_OR_EQUAL, {literal}}, + [&literal](GlobalIndexReader* reader) { return reader->VisitLessOrEqual(literal); }); + } + + Result> VisitGreaterThan(const Literal& literal) override { + return Query( + {QueryOperation::GREATER_THAN, {literal}}, + [&literal](GlobalIndexReader* reader) { return reader->VisitGreaterThan(literal); }); + } + + Result> VisitGreaterOrEqual( + const Literal& literal) override { + return Query( + {QueryOperation::GREATER_OR_EQUAL, {literal}}, + [&literal](GlobalIndexReader* reader) { return reader->VisitGreaterOrEqual(literal); }); + } + + Result> VisitIn( + const std::vector& literals) override { + return Query({QueryOperation::IN, literals}, + [&literals](GlobalIndexReader* reader) { return reader->VisitIn(literals); }); + } + + Result> VisitNotIn( + const std::vector& literals) override { + return Query({QueryOperation::NOT_IN, literals}, [&literals](GlobalIndexReader* reader) { + return reader->VisitNotIn(literals); + }); + } + + Result> VisitStartsWith(const Literal& prefix) override { + return Query({QueryOperation::STARTS_WITH, {prefix}}, [&prefix](GlobalIndexReader* reader) { + return reader->VisitStartsWith(prefix); + }); + } + + Result> VisitEndsWith(const Literal& suffix) override { + return Query({QueryOperation::ENDS_WITH, {suffix}}, [&suffix](GlobalIndexReader* reader) { + return reader->VisitEndsWith(suffix); + }); + } + + Result> VisitContains(const Literal& literal) override { + return Query({QueryOperation::CONTAINS, {literal}}, [&literal](GlobalIndexReader* reader) { + return reader->VisitContains(literal); + }); + } + + Result> VisitLike(const Literal& literal) override { + return Query({QueryOperation::LIKE, {literal}}, + [&literal](GlobalIndexReader* reader) { return reader->VisitLike(literal); }); + } + + Result> VisitVectorSearch( + const std::shared_ptr& vector_search) override { + return Status::Invalid("Primary-key sorted index does not support vector search."); + } + + Result> VisitFullTextSearch( + const std::shared_ptr& full_text_search) override { + return Status::Invalid("Primary-key sorted index does not support full text search."); + } + + bool IsThreadSafe() const override { + return false; + } + + std::string GetIndexType() const override { + return shared_reader_->Group()->Payload()->IndexType(); + } + + private: + Result> Query( + QueryKey key, + const std::function>(GlobalIndexReader*)>& + query) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr group_result, + shared_reader_->Query(key, query)); + return shared_reader_->Localize(group_result, source_index_); + } + + std::shared_ptr shared_reader_; + size_t source_index_; +}; + +Result FindSourceIndex(const PkSortedIndexGroup& group, const DataFileMeta& data_file) { + const std::vector& source_files = group.SourceFiles(); + for (size_t i = 0; i < source_files.size(); i++) { + if (source_files[i].file_name == data_file.file_name && + source_files[i].row_count == data_file.row_count) { + return i; + } + } + return Status::Invalid(fmt::format( + "Data file {} is not covered by its sorted-index source group.", data_file.file_name)); +} +} // namespace + +Result PrimaryKeySortedIndexScan::CreatePlan( + int64_t snapshot_id, const std::vector>& data_splits, + const std::vector& definitions, + const std::vector& index_entries) { + std::unordered_map>> payloads_by_bucket; + for (const IndexManifestEntry& entry : index_entries) { + const std::shared_ptr& payload = entry.index_file; + if (payload == nullptr || !(entry.kind == FileKind::Add())) { + continue; + } + const std::optional& meta = payload->GetGlobalIndexMeta(); + if (meta == std::nullopt || meta.value().source_meta == nullptr) { + continue; + } + payloads_by_bucket[BucketKey(entry.partition, entry.bucket)].push_back(payload); + } + + std::vector scalar_definitions; + for (const PrimaryKeyIndexDefinition& definition : definitions) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE || + definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BITMAP) { + scalar_definitions.push_back(definition); + } + } + + std::unordered_map>> data_files_by_bucket; + for (const std::shared_ptr& split : data_splits) { + if (split == nullptr) { + return Status::Invalid("Primary-key sorted-index scan received a null data split."); + } + if (split->SnapshotId() != snapshot_id) { + return Status::Invalid( + fmt::format("Data split snapshot {} does not match sorted-index scan snapshot {}.", + split->SnapshotId(), snapshot_id)); + } + if (split->IsStreaming()) { + return Status::Invalid("Primary-key sorted-index scan requires batch splits."); + } + if (!split->DeletionFiles().empty() && + split->DeletionFiles().size() != split->DataFiles().size()) { + return Status::Invalid( + "Deletion files must align with data files in a sorted-index split."); + } + std::vector>& data_files = + data_files_by_bucket[BucketKey(split->Partition(), split->Bucket())]; + data_files.insert(data_files.end(), split->DataFiles().begin(), split->DataFiles().end()); + } + + // file name -> field id -> validated group, per bucket. + std::unordered_map< + BucketKey, std::map>>> + groups_by_bucket; + for (const auto& bucket_entry : data_files_by_bucket) { + const BucketKey& bucket = bucket_entry.first; + std::vector> bucket_payloads; + auto payloads_iter = payloads_by_bucket.find(bucket); + if (payloads_iter != payloads_by_bucket.end()) { + bucket_payloads = payloads_iter->second; + } + std::set> active_source_files; + for (const std::shared_ptr& data_file : bucket_entry.second) { + active_source_files.emplace(data_file->file_name, data_file->row_count); + } + std::map>> + groups_by_source; + for (const PrimaryKeyIndexDefinition& definition : scalar_definitions) { + std::vector> definition_payloads; + for (const std::shared_ptr& payload : bucket_payloads) { + const std::optional& meta = payload->GetGlobalIndexMeta(); + if (meta != std::nullopt && definition.IndexType() == payload->IndexType() && + definition.FieldId() == meta.value().index_field_id) { + definition_payloads.push_back(payload); + } + } + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + definition.FieldId(), definition.IndexType(), bucket_entry.second, + definition_payloads); + for (const PkSortedIndexGroup& group : state.Groups()) { + auto shared_group = std::make_shared(group); + for (const PrimaryKeyIndexSourceFile& source_file : group.SourceFiles()) { + if (active_source_files.count({source_file.file_name, source_file.row_count}) == + 0) { + continue; + } + groups_by_source[source_file.file_name][definition.FieldId()] = shared_group; + } + } + } + groups_by_bucket[bucket] = std::move(groups_by_source); + } + + std::vector files; + for (const std::shared_ptr& split : data_splits) { + auto bucket_groups = groups_by_bucket.find(BucketKey(split->Partition(), split->Bucket())); + for (size_t file_index = 0; file_index < split->DataFiles().size(); file_index++) { + const std::shared_ptr& data_file = split->DataFiles()[file_index]; + std::map> groups; + if (bucket_groups != groups_by_bucket.end() && data_file != nullptr) { + auto source_groups = bucket_groups->second.find(data_file->file_name); + if (source_groups != bucket_groups->second.end()) { + groups = source_groups->second; + } + } + files.emplace_back(split, static_cast(file_index), std::move(groups)); + } + } + return Plan(snapshot_id, std::move(files)); +} + +Result PrimaryKeySortedIndexScan::Evaluate( + const Plan& plan, const std::shared_ptr& table_schema, + const std::shared_ptr& predicate, + const std::vector& definitions, + const ReaderFactory& reader_factory) { + std::map definitions_by_field; + for (const PrimaryKeyIndexDefinition& definition : definitions) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE || + definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BITMAP) { + definitions_by_field.emplace(definition.FieldId(), definition); + } + } + + std::unordered_map> + shared_readers; + std::vector files; + files.reserve(plan.Files().size()); + for (const FilePlan& file : plan.Files()) { + GlobalIndexEvaluatorImpl::IndexReadersCreator create_readers = + [&file, &definitions_by_field, &shared_readers, &reader_factory]( + int32_t field_id) -> Result>> { + auto definition_iter = definitions_by_field.find(field_id); + std::shared_ptr group = file.Group(field_id); + if (definition_iter == definitions_by_field.end() || group == nullptr) { + return std::vector>(); + } + auto shared_iter = shared_readers.find(group.get()); + if (shared_iter == shared_readers.end()) { + const PrimaryKeyIndexDefinition& definition = definition_iter->second; + // The shared reader outlives this file plan, so the factory owns a copy of + // the file plan instead of referencing the loop variable. + SharedGroupReader::UnderlyingReaderFactory underlying_factory = + [file_copy = file, definition, group, + &reader_factory]() -> Result> { + return reader_factory(file_copy, definition, *group); + }; + shared_iter = shared_readers + .emplace(group.get(), std::make_shared( + group, std::move(underlying_factory))) + .first; + } + PAIMON_ASSIGN_OR_RAISE(size_t source_index, FindSourceIndex(*group, *file.DataFile())); + std::vector> readers; + readers.push_back( + std::make_shared(shared_iter->second, source_index)); + return readers; + }; + GlobalIndexEvaluatorImpl evaluator(table_schema, create_readers); + Result> result = evaluator.Evaluate(predicate); + if (result.ok()) { + files.emplace_back(file, result.value()); + } else { + // Evaluation failures degrade to a normal scan for this file only. + files.emplace_back(file, nullptr); + } + } + return EvaluatedPlan(plan.SnapshotId(), std::move(files)); +} + +namespace { +class FsGlobalIndexFileReader : public GlobalIndexFileReader { + public: + explicit FsGlobalIndexFileReader(std::shared_ptr file_system) + : file_system_(std::move(file_system)) {} + + Result> GetInputStream( + const std::string& file_path) const override { + return file_system_->Open(file_path); + } + + private: + std::shared_ptr file_system_; +}; +} // namespace + +PrimaryKeySortedIndexScan::ReaderFactory PrimaryKeySortedIndexScan::MakeReaderFactory( + const std::shared_ptr& file_system, + const std::shared_ptr& path_factories, + const std::shared_ptr& table_schema, const std::shared_ptr& pool, + const std::shared_ptr& executor) { + auto file_reader = std::make_shared(file_system); + return [path_factories, table_schema, pool, file_reader, executor]( + const FilePlan& file, const PrimaryKeyIndexDefinition& definition, + const PkSortedIndexGroup& group) -> Result> { + if (definition.GetFamily() != PrimaryKeyIndexDefinition::Family::BTREE) { + // Only the BTree payload reader is wired up; other families keep normal scan + // semantics until their dedicated readers are supported. + return std::shared_ptr(nullptr); + } + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr indexer, + GlobalIndexerFactory::Get(definition.IndexType(), definition.Options())); + if (indexer == nullptr) { + return std::shared_ptr(nullptr); + } + const std::shared_ptr& split = file.SourceSplit(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr path_factory, + path_factories->Get(split->Partition(), split->Bucket())); + const std::shared_ptr& payload = group.Payload(); + const std::optional& payload_meta = payload->GetGlobalIndexMeta(); + if (payload_meta == std::nullopt) { + // Group validation guarantees the metadata; degrade to a normal scan of the + // covered files if it is ever violated, like the Java reader factory. + return Status::Invalid(fmt::format( + "Sorted index payload {} has no global index metadata.", payload->FileName())); + } + std::vector io_metas; + io_metas.emplace_back(path_factory->ToPath(payload), payload->FileSize(), + payload_meta.value().index_meta); + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(definition.FieldId())); + auto arrow_field = DataField::ConvertDataFieldToArrowField(field); + auto arrow_schema = arrow::schema({arrow_field}); + ArrowSchema c_arrow_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema, &c_arrow_schema)); + ScopeGuard guard([&]() { ArrowSchemaRelease(&c_arrow_schema); }); + return indexer->CreateReader(&c_arrow_schema, file_reader, io_metas, pool, executor); + }; +} + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_sorted_index_scan.h b/src/paimon/core/table/source/primary_key_sorted_index_scan.h new file mode 100644 index 000000000..a892f82ce --- /dev/null +++ b/src/paimon/core/table/source/primary_key_sorted_index_scan.h @@ -0,0 +1,188 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/index/pksorted/pk_sorted_index_group.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/table/source/data_split_impl.h" +#include "paimon/core/utils/index_file_path_factories.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/global_index_reader.h" +#include "paimon/global_index/global_index_result.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/predicate.h" +#include "paimon/result.h" + +namespace paimon { +class Executor; + +/// Plans and evaluates source-backed primary-key scalar index groups in file-local +/// row-position space. +/// +/// The scan works on one captured snapshot: data splits and index manifest entries must +/// come from the same snapshot. Every active data file is associated with the validated +/// payload group of its (bucket, field, data level); files without a valid group keep an +/// empty group map and later fall back to a normal scan. Evaluation runs the predicate +/// against the group payloads once per group and query, then localizes group ordinals to +/// per-file physical row positions using the ordered source row-count prefix. +class PrimaryKeySortedIndexScan { + public: + PrimaryKeySortedIndexScan() = delete; + ~PrimaryKeySortedIndexScan() = delete; + + /// One active data file and its complete field-local payload groups. + class FilePlan { + public: + FilePlan(std::shared_ptr source_split, int32_t file_index, + std::map> groups) + : source_split_(std::move(source_split)), + file_index_(file_index), + groups_(std::move(groups)) {} + + const std::shared_ptr& SourceSplit() const { + return source_split_; + } + + int32_t FileIndex() const { + return file_index_; + } + + const std::shared_ptr& DataFile() const { + return source_split_->DataFiles()[file_index_]; + } + + std::shared_ptr Group(int32_t field_id) const { + auto iter = groups_.find(field_id); + return iter == groups_.end() ? nullptr : iter->second; + } + + const std::map>& Groups() const { + return groups_; + } + + private: + std::shared_ptr source_split_; + int32_t file_index_; + std::map> groups_; + }; + + /// Immutable groups for all source files in one captured snapshot. + class Plan { + public: + Plan(int64_t snapshot_id, std::vector files) + : snapshot_id_(snapshot_id), files_(std::move(files)) {} + + int64_t SnapshotId() const { + return snapshot_id_; + } + + const std::vector& Files() const { + return files_; + } + + private: + int64_t snapshot_id_; + std::vector files_; + }; + + /// Optional file-local index result; a null result means that the file requires a + /// normal scan. + class EvaluatedFile { + public: + EvaluatedFile(FilePlan file, std::shared_ptr result) + : file_(std::move(file)), result_(std::move(result)) {} + + const FilePlan& File() const { + return file_; + } + + const std::shared_ptr& IndexResult() const { + return result_; + } + + private: + FilePlan file_; + std::shared_ptr result_; + }; + + /// Predicate results for all source files in one captured snapshot. + class EvaluatedPlan { + public: + EvaluatedPlan(int64_t snapshot_id, std::vector files) + : snapshot_id_(snapshot_id), files_(std::move(files)) {} + + int64_t SnapshotId() const { + return snapshot_id_; + } + + const std::vector& Files() const { + return files_; + } + + private: + int64_t snapshot_id_; + std::vector files_; + }; + + /// Creates a payload reader for one validated group. Returning a null reader marks the + /// index type as unusable so affected predicates keep their normal scan semantics. + using ReaderFactory = std::function>( + const FilePlan& file, const PrimaryKeyIndexDefinition& definition, + const PkSortedIndexGroup& group)>; + + /// Associates every active data file with the validated payload groups of its bucket. + /// + /// `index_entries` must be the ADD entries of the same snapshot carrying global index + /// metadata with source metadata. Splits must be non-streaming and their deletion + /// files, when present, must align with their data files. + static Result CreatePlan(int64_t snapshot_id, + const std::vector>& data_splits, + const std::vector& definitions, + const std::vector& index_entries); + + /// Evaluates the predicate for every planned file. Evaluation failures degrade to a + /// null per-file result instead of failing the scan. + static Result Evaluate(const Plan& plan, + const std::shared_ptr& table_schema, + const std::shared_ptr& predicate, + const std::vector& definitions, + const ReaderFactory& reader_factory); + + /// Creates the default reader factory which opens BTree payloads through the table's + /// index directory layout. Non-BTree families resolve to a null reader and therefore + /// keep normal scan semantics. + static ReaderFactory MakeReaderFactory( + const std::shared_ptr& file_system, + const std::shared_ptr& path_factories, + const std::shared_ptr& table_schema, const std::shared_ptr& pool, + const std::shared_ptr& executor); +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp b/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp new file mode 100644 index 000000000..fdb6ffc7e --- /dev/null +++ b/src/paimon/core/table/source/primary_key_sorted_index_scan_test.cpp @@ -0,0 +1,548 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/table/source/primary_key_sorted_index_scan.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "fmt/format.h" +#include "gtest/gtest.h" +#include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/index/pksorted/pk_sorted_index_file.h" +#include "paimon/core/manifest/file_kind.h" +#include "paimon/core/manifest/file_source.h" +#include "paimon/core/table/source/primary_key_sorted_index_result.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/global_index/global_index_io_meta.h" +#include "paimon/global_index/global_indexer.h" +#include "paimon/global_index/global_indexer_factory.h" +#include "paimon/global_index/io/global_index_file_reader.h" +#include "paimon/global_index/io/global_index_file_writer.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { +constexpr int32_t kPriceFieldId = 1; +constexpr int64_t kSnapshotId = 7; +constexpr int64_t kFileARows = 100; +constexpr int64_t kFileBRows = 200; +constexpr int64_t kTotalRows = kFileARows + kFileBRows; + +class TestGlobalIndexFileWriter : public GlobalIndexFileWriter { + public: + TestGlobalIndexFileWriter(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 fmt::format("{}-index-{}", prefix, 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(std::shared_ptr file_status, + fs_->GetFileStatus(base_path_ + "/" + file_name)); + return 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 TestGlobalIndexFileReader : public GlobalIndexFileReader { + public: + explicit TestGlobalIndexFileReader(const std::shared_ptr& fs) : fs_(fs) {} + + Result> GetInputStream( + const std::string& file_path) const override { + return fs_->Open(file_path); + } + + private: + std::shared_ptr fs_; +}; + +/// A reader stub whose equality result is fully controlled by the test, used to exercise +/// the untrusted-position fallbacks. +class StubGlobalIndexReader : public GlobalIndexReader { + public: + explicit StubGlobalIndexReader(RoaringBitmap64 equal_result) + : equal_result_(std::move(equal_result)) {} + + Result> VisitIsNotNull() override { + return NotEvaluable(); + } + Result> VisitIsNull() override { + return NotEvaluable(); + } + Result> VisitEqual(const Literal& literal) override { + RoaringBitmap64 copy = equal_result_; + return std::make_shared( + [bitmap = std::move(copy)]() -> Result { return bitmap; }); + } + Result> VisitNotEqual(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitLessThan(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitLessOrEqual(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitGreaterThan(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitGreaterOrEqual( + const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitIn( + const std::vector& literals) override { + return NotEvaluable(); + } + Result> VisitNotIn( + const std::vector& literals) override { + return NotEvaluable(); + } + Result> VisitStartsWith(const Literal& prefix) override { + return NotEvaluable(); + } + Result> VisitEndsWith(const Literal& suffix) override { + return NotEvaluable(); + } + Result> VisitContains(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitLike(const Literal& literal) override { + return NotEvaluable(); + } + Result> VisitVectorSearch( + const std::shared_ptr& vector_search) override { + return Status::Invalid("not supported"); + } + Result> VisitFullTextSearch( + const std::shared_ptr& full_text_search) override { + return Status::Invalid("not supported"); + } + bool IsThreadSafe() const override { + return false; + } + std::string GetIndexType() const override { + return "btree"; + } + + private: + static Result> NotEvaluable() { + return std::shared_ptr(nullptr); + } + + RoaringBitmap64 equal_result_; +}; +} // namespace + +class PrimaryKeySortedIndexScanTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + test_dir_ = UniqueTestDirectory::Create("local"); + fs_ = test_dir_->GetFileSystem(); + base_path_ = test_dir_->Str(); + + std::vector fields = { + DataField(0, arrow::field("id", arrow::int64())), + DataField(kPriceFieldId, arrow::field("price", arrow::int64())), + DataField(2, arrow::field("status", arrow::utf8())), + }; + std::map options = {{"pk-btree.index.columns", "price"}}; + table_schema_ = std::make_shared( + /*version=*/3, /*id=*/0, fields, /*highest_field_id=*/2, + /*partition_keys=*/std::vector(), + /*primary_keys=*/std::vector{"id"}, options, + /*comment=*/std::nullopt, /*time_millis=*/0); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*table_schema_)); + definitions_ = definitions.ScalarDefinitions(); + ASSERT_EQ(definitions_.size(), 1); + } + + std::shared_ptr MakeDataFile(const std::string& name, int64_t row_count, + int32_t level, const FileSource& file_source) { + return std::make_shared( + name, /*file_size=*/1024, row_count, + /*min_key=*/BinaryRow::EmptyRow(), /*max_key=*/BinaryRow::EmptyRow(), + /*key_stats=*/SimpleStats::EmptyStats(), /*value_stats=*/SimpleStats::EmptyStats(), + /*min_sequence_number=*/0, /*max_sequence_number=*/row_count, /*schema_id=*/0, level, + /*extra_files=*/std::vector>(), + /*creation_time=*/Timestamp(1721643142456LL, 0), + /*delete_row_count=*/0, /*embedded_index=*/nullptr, file_source, + /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + } + + Result> BuildPayload(std::vector ordinals) { + std::vector source_files = {{"a.parquet", kFileARows}, + {"b.parquet", kFileBRows}}; + arrow::Int64Builder values_builder; + for (int64_t i = 0; i < kTotalRows; i++) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Append(2 * i)); + } + std::shared_ptr sorted_values; + PAIMON_RETURN_NOT_OK_FROM_ARROW(values_builder.Finish(&sorted_values)); + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema_->GetField(kPriceFieldId)); + auto file_writer = std::make_shared(fs_, base_path_); + return PkSortedIndexFile::Build(field, "btree", definitions_[0].Options(), + /*data_level=*/5, source_files, sorted_values, + std::move(ordinals), file_writer, + /*is_external_path=*/false, pool_); + } + + /// Builds the standard payload of this fixture: sources a.parquet(100) + b.parquet(200) + /// on level 5, indexed value at group ordinal `i` is `2 * i`. + Result> BuildPayload() { + std::vector ordinals; + ordinals.reserve(kTotalRows); + for (int64_t i = 0; i < kTotalRows; i++) { + ordinals.push_back(i); + } + return BuildPayload(std::move(ordinals)); + } + + std::shared_ptr MakeSplit( + const std::vector>& files, bool raw_convertible, + const std::vector>& deletion_files = {}) { + std::vector> data_files = files; + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, + base_path_ + "/bucket-0", std::move(data_files)); + builder.WithSnapshot(kSnapshotId).IsStreaming(false).RawConvertible(raw_convertible); + if (!deletion_files.empty()) { + builder.WithDataDeletionFiles(deletion_files); + } + EXPECT_OK_AND_ASSIGN(std::shared_ptr split, builder.Build()); + return split; + } + + std::vector MakeEntries(const std::shared_ptr& payload) { + return {IndexManifestEntry(FileKind::Add(), BinaryRow::EmptyRow(), /*bucket=*/0, payload)}; + } + + PrimaryKeySortedIndexScan::ReaderFactory PayloadReaderFactory() { + std::shared_ptr fs = fs_; + std::string base_path = base_path_; + std::shared_ptr table_schema = table_schema_; + std::shared_ptr pool = pool_; + return [fs, base_path, table_schema, pool]( + const PrimaryKeySortedIndexScan::FilePlan& file, + const PrimaryKeyIndexDefinition& definition, + const PkSortedIndexGroup& group) -> Result> { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr indexer, + GlobalIndexerFactory::Get(definition.IndexType(), definition.Options())); + if (indexer == nullptr) { + return Status::Invalid("btree indexer is not registered"); + } + const std::shared_ptr& payload = group.Payload(); + std::vector io_metas; + io_metas.emplace_back(base_path + "/" + payload->FileName(), payload->FileSize(), + payload->GetGlobalIndexMeta().value().index_meta); + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(definition.FieldId())); + auto arrow_field = DataField::ConvertDataFieldToArrowField(field); + auto arrow_schema = arrow::schema({arrow_field}); + ArrowSchema c_arrow_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema, &c_arrow_schema)); + auto file_reader = std::make_shared(fs); + return indexer->CreateReader(&c_arrow_schema, file_reader, io_metas, pool); + }; + } + + Result>> PlanEvaluateConvert( + const std::vector>& splits, + const std::vector& entries, const std::shared_ptr& predicate, + const PrimaryKeySortedIndexScan::ReaderFactory& reader_factory) { + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeySortedIndexScan::Plan plan, + PrimaryKeySortedIndexScan::CreatePlan(kSnapshotId, splits, definitions_, entries)); + PAIMON_ASSIGN_OR_RAISE(PrimaryKeySortedIndexScan::EvaluatedPlan evaluated, + PrimaryKeySortedIndexScan::Evaluate(plan, table_schema_, predicate, + definitions_, reader_factory)); + return PrimaryKeySortedIndexResult::ToSplits(evaluated); + } + + std::shared_ptr PriceEqual(int64_t value) { + return PredicateBuilder::Equal(/*field_index=*/1, "price", FieldType::BIGINT, + Literal(value)); + } + + std::shared_ptr pool_; + std::shared_ptr test_dir_; + std::shared_ptr fs_; + std::string base_path_; + std::shared_ptr table_schema_; + std::vector definitions_; +}; + +TEST_F(PrimaryKeySortedIndexScanTest, EqualNarrowsToSingleFileRange) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + // Value 10 sits at group ordinal 5, i.e. row 5 of a.parquet. + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10), PayloadReaderFactory())); + ASSERT_EQ(splits.size(), 1); + auto indexed_split = std::dynamic_pointer_cast(splits[0]); + ASSERT_TRUE(indexed_split != nullptr); + auto inner_split = std::dynamic_pointer_cast(indexed_split->GetDataSplit()); + ASSERT_TRUE(inner_split != nullptr); + ASSERT_EQ(inner_split->DataFiles().size(), 1); + ASSERT_EQ(inner_split->DataFiles()[0]->file_name, "a.parquet"); + ASSERT_EQ(indexed_split->RowRanges().size(), 1); + ASSERT_EQ(indexed_split->RowRanges()[0].from, 5); + ASSERT_EQ(indexed_split->RowRanges()[0].to, 5); +} + +TEST_F(PrimaryKeySortedIndexScanTest, BuildRejectsDuplicateOrdinals) { + std::vector ordinals; + ordinals.reserve(kTotalRows); + for (int64_t i = 0; i < kTotalRows; i++) { + ordinals.push_back(i); + } + ordinals[1] = 0; + ASSERT_NOK_WITH_MSG(BuildPayload(std::move(ordinals)).status(), + "Row id 0 appears more than once"); +} + +TEST_F(PrimaryKeySortedIndexScanTest, RangeSpansFileBoundary) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + // Values in [190, 210] sit at group ordinals 95..105: rows 95..99 of a.parquet and + // rows 0..5 of b.parquet. + std::shared_ptr lower = PredicateBuilder::GreaterOrEqual( + /*field_index=*/1, "price", FieldType::BIGINT, Literal(static_cast(190))); + std::shared_ptr upper = PredicateBuilder::LessOrEqual( + /*field_index=*/1, "price", FieldType::BIGINT, Literal(static_cast(210))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate, + PredicateBuilder::And({lower, upper})); + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), predicate, PayloadReaderFactory())); + ASSERT_EQ(splits.size(), 2); + auto indexed_a = std::dynamic_pointer_cast(splits[0]); + auto indexed_b = std::dynamic_pointer_cast(splits[1]); + ASSERT_TRUE(indexed_a != nullptr); + ASSERT_TRUE(indexed_b != nullptr); + ASSERT_EQ(indexed_a->RowRanges().size(), 1); + ASSERT_EQ(indexed_a->RowRanges()[0].from, 95); + ASSERT_EQ(indexed_a->RowRanges()[0].to, 99); + ASSERT_EQ(indexed_b->RowRanges().size(), 1); + ASSERT_EQ(indexed_b->RowRanges()[0].from, 0); + ASSERT_EQ(indexed_b->RowRanges()[0].to, 5); +} + +TEST_F(PrimaryKeySortedIndexScanTest, EmptyResultOmitsAllFiles) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + // All indexed values are even, so 11 matches nothing. + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(11), PayloadReaderFactory())); + ASSERT_TRUE(splits.empty()); +} + +TEST_F(PrimaryKeySortedIndexScanTest, UnindexedFieldPredicateFallsBack) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/2, "status", FieldType::STRING, Literal(FieldType::STRING, "hit", 3)); + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), predicate, PayloadReaderFactory())); + ASSERT_EQ(1, splits.size()); + ASSERT_EQ(split, splits[0]); +} + +TEST_F(PrimaryKeySortedIndexScanTest, UncoveredFileFallsBackOthersNarrow) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact()), + MakeDataFile("c.parquet", 50, 0, FileSource::Append())}, + /*raw_convertible=*/true); + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10), PayloadReaderFactory())); + // a.parquet narrows to an indexed split, b.parquet is omitted, c.parquet has no + // coverage and keeps a normal single-file scan. + ASSERT_EQ(splits.size(), 2); + auto indexed_split = std::dynamic_pointer_cast(splits[0]); + ASSERT_TRUE(indexed_split != nullptr); + auto fallback_split = std::dynamic_pointer_cast(splits[1]); + ASSERT_TRUE(fallback_split != nullptr); + ASSERT_EQ(fallback_split->DataFiles().size(), 1); + ASSERT_EQ(fallback_split->DataFiles()[0]->file_name, "c.parquet"); +} + +TEST_F(PrimaryKeySortedIndexScanTest, NonRawConvertibleSplitPreserved) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/false); + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10), PayloadReaderFactory())); + ASSERT_EQ(splits.size(), 1); + ASSERT_EQ(splits[0].get(), split.get()); +} + +TEST_F(PrimaryKeySortedIndexScanTest, InvalidRowRangePayloadFallsBack) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + // Rebuild the payload metadata with a row range end beyond the source rows: the group + // validation must reject it and every file keeps a normal scan. + const GlobalIndexMeta& meta = payload->GetGlobalIndexMeta().value(); + auto broken_payload = std::make_shared( + payload->IndexType(), payload->FileName(), payload->FileSize(), payload->RowCount(), + std::nullopt, std::nullopt, + GlobalIndexMeta(meta.row_range_start, meta.row_range_end + 1, meta.index_field_id, + meta.extra_field_ids, meta.index_meta, meta.source_meta)); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + ASSERT_OK_AND_ASSIGN(std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(broken_payload), PriceEqual(10), + PayloadReaderFactory())); + ASSERT_EQ(1, splits.size()); + ASSERT_EQ(split, splits[0]); +} + +TEST_F(PrimaryKeySortedIndexScanTest, OutOfRangePositionsFailAllCoveredFiles) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::shared_ptr split = + MakeSplit({MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true); + RoaringBitmap64 poisoned; + poisoned.Add(5); + poisoned.Add(kTotalRows + 10); + PrimaryKeySortedIndexScan::ReaderFactory stub_factory = + [&poisoned](const PrimaryKeySortedIndexScan::FilePlan& file, + const PrimaryKeyIndexDefinition& definition, + const PkSortedIndexGroup& group) -> Result> { + return std::make_shared(poisoned); + }; + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10), stub_factory)); + // Both covered files fall back together, preserving the planner's original bin packing. + ASSERT_EQ(1, splits.size()); + ASSERT_EQ(split, splits[0]); +} + +TEST_F(PrimaryKeySortedIndexScanTest, OverFragmentedResultFallsBack) { + // One data file, 20000 rows; every second row selected produces > 4096 ranges. + std::vector source_files = {{"big.parquet", 20000}}; + std::shared_ptr split = MakeSplit( + {MakeDataFile("big.parquet", 20000, 5, FileSource::Compact())}, /*raw_convertible=*/true); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr source_meta_bytes, ([&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create(5, source_files)); + return source_meta.Serialize(pool_); + }())); + auto big_payload = std::make_shared( + "btree", "big-index-file", /*file_size=*/1, /*row_count=*/20000, std::nullopt, std::nullopt, + GlobalIndexMeta(0, 19999, kPriceFieldId, std::nullopt, nullptr, source_meta_bytes)); + RoaringBitmap64 fragmented; + for (int64_t i = 0; i < 20000; i += 2) { + fragmented.Add(i); + } + PrimaryKeySortedIndexScan::ReaderFactory stub_factory = + [&fragmented]( + const PrimaryKeySortedIndexScan::FilePlan& file, + const PrimaryKeyIndexDefinition& definition, + const PkSortedIndexGroup& group) -> Result> { + return std::make_shared(fragmented); + }; + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(big_payload), PriceEqual(10), stub_factory)); + ASSERT_EQ(splits.size(), 1); + ASSERT_TRUE(std::dynamic_pointer_cast(splits[0]) == nullptr); +} + +TEST_F(PrimaryKeySortedIndexScanTest, DeletionFileStaysAlignedWithIndexedFile) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + DeletionFile deletion_file("dv-a", /*offset=*/0, /*length=*/16, /*cardinality=*/1); + std::shared_ptr split = MakeSplit( + {MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact()), + MakeDataFile("b.parquet", kFileBRows, 5, FileSource::Compact())}, + /*raw_convertible=*/true, {std::optional(deletion_file), std::nullopt}); + ASSERT_OK_AND_ASSIGN( + std::vector> splits, + PlanEvaluateConvert({split}, MakeEntries(payload), PriceEqual(10), PayloadReaderFactory())); + ASSERT_EQ(splits.size(), 1); + auto indexed_split = std::dynamic_pointer_cast(splits[0]); + ASSERT_TRUE(indexed_split != nullptr); + auto inner_split = std::dynamic_pointer_cast(indexed_split->GetDataSplit()); + ASSERT_TRUE(inner_split != nullptr); + ASSERT_EQ(inner_split->DeletionFiles().size(), 1); + ASSERT_TRUE(inner_split->DeletionFiles()[0] != std::nullopt); + ASSERT_EQ(inner_split->DeletionFiles()[0].value().path, "dv-a"); +} + +TEST_F(PrimaryKeySortedIndexScanTest, SnapshotMismatchIsRejected) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload, BuildPayload()); + std::vector> files = { + MakeDataFile("a.parquet", kFileARows, 5, FileSource::Compact())}; + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, base_path_, + std::move(files)); + builder.WithSnapshot(kSnapshotId + 1).IsStreaming(false).RawConvertible(true); + ASSERT_OK_AND_ASSIGN(std::shared_ptr split, builder.Build()); + Result plan = PrimaryKeySortedIndexScan::CreatePlan( + kSnapshotId, {split}, definitions_, MakeEntries(payload)); + ASSERT_NOK(plan.status()); +} + +} // namespace paimon::test diff --git a/src/paimon/core/table/source/snapshot/snapshot_reader.h b/src/paimon/core/table/source/snapshot/snapshot_reader.h index 315ca6e07..c425fd863 100644 --- a/src/paimon/core/table/source/snapshot/snapshot_reader.h +++ b/src/paimon/core/table/source/snapshot/snapshot_reader.h @@ -92,6 +92,10 @@ class SnapshotReader { return scan_->GetSnapshotManager(); } + const std::unique_ptr& GetIndexFileHandler() const { + return index_file_handler_; + } + std::shared_ptr GetNonPartitionPredicate() const { return scan_->GetNonPartitionPredicate(); } diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 62bda904d..fa8b0ccd2 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -33,6 +33,7 @@ #include "paimon/common/utils/options_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/index/index_file_handler.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" #include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" @@ -51,6 +52,7 @@ #include "paimon/core/table/source/data_table_batch_scan.h" #include "paimon/core/table/source/data_table_stream_scan.h" #include "paimon/core/table/source/merge_tree_split_generator.h" +#include "paimon/core/table/source/primary_key_index_batch_scan.h" #include "paimon/core/table/source/read_optimized_scan_options.h" #include "paimon/core/table/source/realtime_table_scan.h" #include "paimon/core/table/source/snapshot/snapshot_reader.h" @@ -339,12 +341,22 @@ Result> NewDataTableScan(const std::shared_ptrGetSnapshotManager(), core_options.GetFileSystem(), context->GetScanFilters()); } - if (!core_options.DataEvolutionEnabled()) { - return batch_scan; + if (core_options.DataEvolutionEnabled()) { + return std::make_unique( + context->GetPath(), snapshot_reader, std::move(batch_scan), + context->GetGlobalIndexResult(), core_options, context->GetMemoryPool(), + context->GetExecutor()); + } + if (pk_table && !read_optimized && core_options.GlobalIndexEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(*table_schema)); + if (!definitions.ScalarDefinitions().empty()) { + return PrimaryKeyIndexBatchScan::Create(snapshot_reader, std::move(batch_scan), + table_schema, path_factory, core_options, + context->GetMemoryPool()); + } } - return std::make_unique( - context->GetPath(), snapshot_reader, std::move(batch_scan), context->GetGlobalIndexResult(), - core_options, context->GetMemoryPool(), context->GetExecutor()); + return batch_scan; } } // namespace