diff --git a/docs/source/user_guide/primary_key_global_index.rst b/docs/source/user_guide/primary_key_global_index.rst index dcba88fb9..9307d938d 100644 --- a/docs/source/user_guide/primary_key_global_index.rst +++ b/docs/source/user_guide/primary_key_global_index.rst @@ -24,11 +24,11 @@ Paimon 2.0 primary-key tables support *source-backed* global scalar indexes covers an immutable ordered source group from 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. +paimon-cpp supports the BTree build, maintenance, and read lifecycle of this protocol for +fixed-bucket primary-key tables. Compaction automatically builds one immutable payload for +each indexed field and positive data level. Ordinary batch scans evaluate the indexed part +of the predicate against validated payload groups and narrow covered files to indexed +splits carrying file-local row ranges. No dedicated build or query API is required. Table requirements ------------------ @@ -37,7 +37,8 @@ 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 +- fixed bucket (``bucket > 0``) for automatic C++ maintenance; Java-compatible postpone + bucket schemas remain readable, but the C++ postpone writer does not build payloads - ``'deletion-vectors.enabled' = 'true'`` and ``'deletion-vectors.merge-on-read' = 'false'`` Semantics @@ -50,6 +51,23 @@ Semantics level, an active source's row count differs, its metadata or row range is invalid, or another payload exists for that level. Active files without accepted coverage are scanned normally. +- Index construction reads every physical source row without applying deletion vectors, + orders source files by file name, and externally sorts ``(value, group row id)``. During + maintenance, missing, duplicate, malformed, or incomplete payloads cause their complete current + level to be rebuilt. A payload that still covers every active source is reused even if it also + lists retired sources. Data files and the corresponding index ADD / DELETE entries are committed + in the same snapshot. +- The builder uses the existing write-buffer and spill settings. A write context needs a + temporary directory when a level exceeds the in-memory write buffer and spill is enabled. +- If payload construction fails, the data-file transition is still committed. Uncovered files at + that level fall back to scanning while any previously usable payload remains active, and a later + maintenance attempt can rebuild the complete current source group. Structural commit-increment + errors are still rejected. +- Snapshot expiration retains payloads referenced by the snapshots in its retention set and + current-branch live tags, and removes retired payloads before their index manifests, including + payloads on an external index path. Expiration is rejected while another branch exists until + cross-branch file retention is supported. Orphan cleanup covers table-local index manifests and + payloads; it does not enumerate a potentially shared global-index external path. - ``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 @@ -70,6 +88,6 @@ Current scope 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. +- Automatic maintenance is synchronous during prepare-commit. Java's asynchronous build + scheduling, retries, fairness metrics, and manual rebuild actions are not part of the C++ + API. Realtime and postpone-bucket writers do not build source-backed payloads. diff --git a/include/paimon/defs.h b/include/paimon/defs.h index e05faefd8..66b1af668 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h @@ -377,6 +377,10 @@ struct PAIMON_EXPORT Options { /// a table is never compacted. static const char DELETION_VECTORS_ENABLED[]; + /// "pk-clustering-override" - Whether primary-key clustering columns override the primary + /// keys when clustering data. Default value is false. + static const char PK_CLUSTERING_OVERRIDE[]; + /// "deletion-vector.index-file.target-size" - The target size of deletion vector index file. /// Default value is 2MB. static const char DELETION_VECTOR_INDEX_FILE_TARGET_SIZE[]; diff --git a/include/paimon/orphan_files_cleaner.h b/include/paimon/orphan_files_cleaner.h index c2349a250..2c1847002 100644 --- a/include/paimon/orphan_files_cleaner.h +++ b/include/paimon/orphan_files_cleaner.h @@ -165,9 +165,9 @@ class PAIMON_EXPORT CleanContextBuilder { /// by Paimon C++, we implemented a strong pattern-matching validation, deleting only files in /// patterns we recognize. /// -/// @note `OrphanFilesCleaner` in Paimon C++ only support cleaning append table, do not support -/// cleaning table with tag, table with external paths, table with branch, table with index, table -/// with changelog, and primary key table. +/// @note `OrphanFilesCleaner` in Paimon C++ does not support cleaning tables with tags, branches, +/// external data paths, or changelog manifests. Global-index external paths are not enumerated; +/// snapshot expiration owns deletion of external index payloads. class PAIMON_EXPORT OrphanFilesCleaner { public: virtual ~OrphanFilesCleaner() = default; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 051eba324..669d6b2ce 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -263,8 +263,11 @@ set(PAIMON_CORE_SRCS 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/pk/bucketed_primary_key_index_maintainer.cpp core/index/pksorted/pk_sorted_index_group.cpp core/index/pksorted/pk_sorted_bucket_index_state.cpp + core/index/pksorted/pk_sorted_data_file_reader.cpp + core/index/pksorted/pk_sorted_index_builder.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 @@ -757,6 +760,7 @@ if(PAIMON_BUILD_TESTS) 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/pk/bucketed_primary_key_index_maintainer_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 diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index f0b0f5611..996813c8b 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp @@ -92,6 +92,7 @@ const char Options::FALLBACK_DEDUPLICATE_IGNORE_DELETE[] = "deduplicate.ignore-d const char Options::FALLBACK_PARTIAL_UPDATE_IGNORE_DELETE[] = "partial-update.ignore-delete"; const char Options::FIELDS_DEFAULT_AGG_FUNC[] = "fields.default-aggregate-function"; const char Options::DELETION_VECTORS_ENABLED[] = "deletion-vectors.enabled"; +const char Options::PK_CLUSTERING_OVERRIDE[] = "pk-clustering-override"; const char Options::DELETION_VECTOR_INDEX_FILE_TARGET_SIZE[] = "deletion-vector.index-file.target-size"; const char Options::DELETION_VECTOR_BITMAP64[] = "deletion-vectors.bitmap64"; diff --git a/src/paimon/core/index/index_file_handler.cpp b/src/paimon/core/index/index_file_handler.cpp index 9b45e748a..0fd776671 100644 --- a/src/paimon/core/index/index_file_handler.cpp +++ b/src/paimon/core/index/index_file_handler.cpp @@ -19,6 +19,7 @@ #include "paimon/core/index/index_file_handler.h" +#include #include #include @@ -26,6 +27,22 @@ #include "paimon/status.h" namespace paimon { +namespace { + +constexpr char kDataEvolutionSourceMetaMagic[] = "DEIX"; + +} // namespace + +bool IndexFileHandler::IsPrimaryKeySourceIndex(const IndexFileMeta& index_file) { + const std::optional& global_index_meta = index_file.GetGlobalIndexMeta(); + if (!global_index_meta.has_value() || global_index_meta->source_meta == nullptr) { + return false; + } + const std::shared_ptr& source_meta = global_index_meta->source_meta; + constexpr size_t kMagicSize = sizeof(kDataEvolutionSourceMetaMagic) - 1; + return source_meta->size() < kMagicSize || + std::memcmp(source_meta->data(), kDataEvolutionSourceMetaMagic, kMagicSize) != 0; +} Result IndexFileHandler::Scan( const Snapshot& snapshot, const std::string& index_type, @@ -73,4 +90,20 @@ Result>> IndexFileHandler::Scan( return std::vector>{}; } +Result>> IndexFileHandler::ScanPrimaryKeyIndexes( + const Snapshot& snapshot, const BinaryRow& partition, int32_t bucket) const { + std::function(const IndexManifestEntry&)> filter = + [&partition, bucket](const IndexManifestEntry& entry) -> bool { + return entry.partition == partition && entry.bucket == bucket && + IsPrimaryKeySourceIndex(*entry.index_file); + }; + PAIMON_ASSIGN_OR_RAISE(std::vector entries, Scan(snapshot, filter)); + std::vector> result; + result.reserve(entries.size()); + for (const IndexManifestEntry& entry : entries) { + result.push_back(entry.index_file); + } + return result; +} + } // namespace paimon diff --git a/src/paimon/core/index/index_file_handler.h b/src/paimon/core/index/index_file_handler.h index a84840a45..06cb00253 100644 --- a/src/paimon/core/index/index_file_handler.h +++ b/src/paimon/core/index/index_file_handler.h @@ -55,6 +55,10 @@ class IndexFileHandler { dv_bitmap64_(dv_bitmap64), pool_(pool) {} + /// Returns whether an index file carries primary-key source metadata rather than Java's + /// data-evolution source metadata. + static bool IsPrimaryKeySourceIndex(const IndexFileMeta& index_file); + /// 1.Scan specified index_type index. 2.Cluster with partition & bucket. Result Scan(const Snapshot& snapshot, const std::string& index_type, const std::unordered_set& partitions) const; @@ -64,6 +68,10 @@ class IndexFileHandler { const BinaryRow& partition, int32_t bucket) const; + /// Scan primary-key source-backed index payloads for a partition and bucket. + Result>> ScanPrimaryKeyIndexes( + const Snapshot& snapshot, const BinaryRow& partition, int32_t bucket) const; + /// Scan specified all typed index. Result> Scan( const Snapshot& snapshot, diff --git a/src/paimon/core/index/index_file_handler_test.cpp b/src/paimon/core/index/index_file_handler_test.cpp index a23591f94..dfc0be8cb 100644 --- a/src/paimon/core/index/index_file_handler_test.cpp +++ b/src/paimon/core/index/index_file_handler_test.cpp @@ -31,6 +31,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/deletion_vectors_index_file.h" #include "paimon/core/index/deletion_vector_meta.h" +#include "paimon/core/index/global_index_meta.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" @@ -288,6 +289,31 @@ TEST_F(IndexFileHandlerTest, TestScanWithNoIndexManifest) { ASSERT_TRUE(index_entries.empty()); } +TEST_F(IndexFileHandlerTest, TestScanPrimaryKeyIndexesBySourceMetadata) { + std::string table_path = + paimon::test::GetDataDir() + "/orc/pk_btree_source_meta.db/pk_btree_source_meta/"; + ASSERT_OK_AND_ASSIGN(CoreOptions core_options, + CoreOptions::FromMap({{Options::MANIFEST_FORMAT, "orc"}})); + ASSERT_OK_AND_ASSIGN(std::unique_ptr index_file_handler, + CreateIndexFileHandler(table_path, core_options)); + + SnapshotManager snapshot_manager(core_options.GetFileSystem(), table_path); + ASSERT_OK_AND_ASSIGN(Snapshot snapshot, snapshot_manager.LoadSnapshot(/*snapshot_id=*/5)); + ASSERT_OK_AND_ASSIGN( + std::vector> source_indexes, + index_file_handler->ScanPrimaryKeyIndexes(snapshot, BinaryRow::EmptyRow(), /*bucket=*/0)); + + ASSERT_EQ(source_indexes.size(), 1); + const std::optional& global_index_meta = + source_indexes[0]->GetGlobalIndexMeta(); + ASSERT_TRUE(global_index_meta.has_value()); + ASSERT_NE(global_index_meta->source_meta, nullptr); + + ASSERT_OK_AND_ASSIGN(source_indexes, index_file_handler->ScanPrimaryKeyIndexes( + snapshot, BinaryRow::EmptyRow(), /*bucket=*/1)); + ASSERT_TRUE(source_indexes.empty()); +} + TEST_F(IndexFileHandlerTest, TestScanByPartitionBucketAndReadAllDeletionVectors) { std::string table_path = paimon::test::GetDataDir() + "/orc/pk_table_with_dv_cardinality.db/pk_table_with_dv_cardinality/"; diff --git a/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.cpp b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.cpp new file mode 100644 index 000000000..f32166466 --- /dev/null +++ b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.cpp @@ -0,0 +1,294 @@ +/* + * 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/bucketed_primary_key_index_maintainer.h" + +#include +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/global_index/btree/btree_defs.h" +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/index_file_handler.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pk/primary_key_index_source_policy.h" +#include "paimon/core/index/pksorted/pk_sorted_bucket_index_state.h" +#include "paimon/core/index/pksorted/pk_sorted_index_builder.h" +#include "paimon/core/index/pksorted/pk_sorted_index_group.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/logging.h" + +namespace paimon { +namespace { + +Logger* GetLogger() { + static std::unique_ptr logger = Logger::GetLogger("BucketedPrimaryKeyIndexMaintainer"); + return logger.get(); +} + +void RemoveDataFiles(const std::vector>& files, + std::map>* active) { + for (const std::shared_ptr& file : files) { + if (file != nullptr) { + active->erase(file->file_name); + } + } +} + +Status AddSourceFiles(const std::vector>& files, + std::map>* active) { + for (const std::shared_ptr& file : files) { + if (file == nullptr) { + return Status::Invalid("Primary-key index data increment contains a null file."); + } + if (PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + (*active)[file->file_name] = file; + } + } + return Status::OK(); +} + +Status ValidateAppendFiles(const std::vector>& files) { + for (const std::shared_ptr& file : files) { + if (file == nullptr) { + return Status::Invalid("Primary-key index append increment contains a null file."); + } + if (PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + return Status::Invalid(fmt::format( + "Append file {} must not be a primary-key sorted-index source.", file->file_name)); + } + } + return Status::OK(); +} + +std::string PayloadIdentity(const std::shared_ptr& payload) { + if (payload == nullptr) { + return std::string(); + } + return payload->ExternalPath().value_or(payload->FileName()); +} + +void AddUniquePayload(const std::shared_ptr& payload, + std::unordered_set* identities, + std::vector>* payloads) { + std::string identity = PayloadIdentity(payload); + if (!identity.empty() && identities->insert(identity).second) { + payloads->push_back(payload); + } +} + +bool IsPrimaryKeyBTreePayload(const std::shared_ptr& payload) { + return payload != nullptr && payload->IndexType() == BtreeDefs::kIdentifier && + IndexFileHandler::IsPrimaryKeySourceIndex(*payload); +} + +bool CoversAllSources(const std::vector& group_sources, + const std::vector& desired_sources) { + size_t group_index = 0; + for (const PrimaryKeyIndexSourceFile& desired : desired_sources) { + while (group_index < group_sources.size() && + group_sources[group_index].file_name < desired.file_name) { + group_index++; + } + if (group_index == group_sources.size() || group_sources[group_index] != desired) { + return false; + } + group_index++; + } + return true; +} + +} // namespace + +Result> +BucketedPrimaryKeyIndexMaintainer::Factory::Create( + const std::string& root_path, const std::string& branch, + const std::shared_ptr& table_schema, + const std::vector& definitions, + const std::shared_ptr& path_factory, + const std::shared_ptr& index_file_handler, const CoreOptions& options, + const std::shared_ptr& io_manager, bool enable_multi_thread_spill, + const std::shared_ptr& executor, const std::shared_ptr& pool) { + std::vector btree_definitions; + for (const PrimaryKeyIndexDefinition& definition : definitions) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE) { + btree_definitions.push_back(definition); + } + } + std::sort(btree_definitions.begin(), btree_definitions.end(), + [](const PrimaryKeyIndexDefinition& left, const PrimaryKeyIndexDefinition& right) { + return left.FieldId() < right.FieldId(); + }); + return std::shared_ptr(new Factory( + root_path, branch, table_schema, std::move(btree_definitions), path_factory, + index_file_handler, options, io_manager, enable_multi_thread_spill, executor, pool)); +} + +Result> +BucketedPrimaryKeyIndexMaintainer::Factory::CreateMaintainer( + const BinaryRow& partition, int32_t bucket, + const std::vector>& restored_data_files, + const std::vector>& restored_payloads) const { + std::map> active_data_files; + PAIMON_RETURN_NOT_OK(AddSourceFiles(restored_data_files, &active_data_files)); + std::vector fields; + fields.reserve(definitions_.size()); + for (const PrimaryKeyIndexDefinition& definition : definitions_) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr builder, + PkSortedIndexBuilder::Create(root_path_, branch_, partition, bucket, table_schema_, + definition, path_factory_, options_, io_manager_, + enable_multi_thread_spill_, executor_, pool_)); + fields.push_back( + FieldMaintainer{definition, std::shared_ptr(std::move(builder))}); + } + return std::shared_ptr(new BucketedPrimaryKeyIndexMaintainer( + std::move(fields), std::move(active_data_files), restored_payloads)); +} + +Status BucketedPrimaryKeyIndexMaintainer::PrepareCommit(CommitIncrement* increment) { + if (increment == nullptr) { + return Status::Invalid("Primary-key index commit increment is null."); + } + auto previous_data_files = active_data_files_; + const DataIncrement& data_increment = increment->GetNewFilesIncrement(); + const CompactIncrement& compact_increment = increment->GetCompactIncrement(); + PAIMON_RETURN_NOT_OK(ValidateAppendFiles(data_increment.NewFiles())); + RemoveDataFiles(compact_increment.CompactBefore(), &active_data_files_); + Status update_status = AddSourceFiles(compact_increment.CompactAfter(), &active_data_files_); + if (!update_status.ok()) { + active_data_files_ = std::move(previous_data_files); + return update_status; + } + + std::vector> active_data; + active_data.reserve(active_data_files_.size()); + for (const auto& file : active_data_files_) { + active_data.push_back(file.second); + } + + std::vector> deleted_payloads; + std::vector> new_payloads; + std::unordered_set deleted_identities; + std::unordered_set new_identities; + + std::set owned_btree_field_ids; + for (const FieldMaintainer& field : fields_) { + owned_btree_field_ids.insert(field.definition.FieldId()); + } + for (const std::shared_ptr& payload : active_payloads_) { + if (!IsPrimaryKeyBTreePayload(payload)) { + continue; + } + const std::optional& meta = payload->GetGlobalIndexMeta(); + if (meta != std::nullopt && meta->source_meta != nullptr && + owned_btree_field_ids.count(meta->index_field_id) == 0) { + AddUniquePayload(payload, &deleted_identities, &deleted_payloads); + } + } + + for (const FieldMaintainer& field : fields_) { + std::vector> field_payloads; + for (const std::shared_ptr& payload : active_payloads_) { + if (!IsPrimaryKeyBTreePayload(payload) || + payload->IndexType() != field.definition.IndexType()) { + continue; + } + const std::optional& meta = payload->GetGlobalIndexMeta(); + if (meta != std::nullopt && meta->index_field_id == field.definition.FieldId()) { + field_payloads.push_back(payload); + } + } + PkSortedBucketIndexState state = PkSortedBucketIndexState::FromActiveDataFiles( + field.definition.FieldId(), field.definition.IndexType(), active_data, field_payloads); + std::map> current_groups_by_level; + for (const std::shared_ptr& group : state.Groups()) { + current_groups_by_level.emplace(group->DataLevel(), group); + } + for (const std::shared_ptr& rejected : state.RejectedPayloads()) { + AddUniquePayload(rejected, &deleted_identities, &deleted_payloads); + } + + std::map>> desired_by_level; + for (const std::shared_ptr& file : active_data) { + if (file != nullptr && PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + desired_by_level[file->level].push_back(file); + } + } + for (auto& level_files : desired_by_level) { + std::sort(level_files.second.begin(), level_files.second.end(), + [](const std::shared_ptr& left, + const std::shared_ptr& right) { + return left->file_name < right->file_name; + }); + std::vector desired_sources; + desired_sources.reserve(level_files.second.size()); + for (const std::shared_ptr& file : level_files.second) { + desired_sources.emplace_back(file->file_name, file->row_count); + } + auto current_group = current_groups_by_level.find(level_files.first); + if (current_group != current_groups_by_level.end() && + CoversAllSources(current_group->second->SourceFiles(), desired_sources)) { + continue; + } + Result> build_result = + field.builder->Build(level_files.second); + if (!build_result.ok()) { + PAIMON_LOG_WARN( + GetLogger(), + "Failed to build primary-key BTree index for column %s at data level %d; " + "leaving uncovered files on normal scan fallback. %s", + field.definition.Column().c_str(), level_files.first, + build_result.status().ToString().c_str()); + continue; + } + if (current_group != current_groups_by_level.end()) { + AddUniquePayload(current_group->second->Payload(), &deleted_identities, + &deleted_payloads); + } + AddUniquePayload(std::move(build_result).value(), &new_identities, &new_payloads); + } + } + + std::vector> next_payloads; + next_payloads.reserve(active_payloads_.size() + new_payloads.size()); + for (const std::shared_ptr& payload : active_payloads_) { + if (deleted_identities.count(PayloadIdentity(payload)) == 0) { + next_payloads.push_back(payload); + } + } + next_payloads.insert(next_payloads.end(), new_payloads.begin(), new_payloads.end()); + active_payloads_ = std::move(next_payloads); + + bool has_compaction_transition = + !compact_increment.CompactBefore().empty() || !compact_increment.CompactAfter().empty(); + if (has_compaction_transition) { + increment->GetCompactIncrement().AddNewIndexFiles(std::move(new_payloads)); + increment->GetCompactIncrement().AddDeletedIndexFiles(std::move(deleted_payloads)); + } else { + increment->GetNewFilesIncrement().AddNewIndexFiles(std::move(new_payloads)); + increment->GetNewFilesIncrement().AddDeletedIndexFiles(std::move(deleted_payloads)); + } + return Status::OK(); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.h b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.h new file mode 100644 index 000000000..cb3e00562 --- /dev/null +++ b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer.h @@ -0,0 +1,132 @@ +/* + * 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/common/data/binary_row.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/result.h" + +namespace paimon { +class CommitIncrement; +class Executor; +class FileStorePathFactory; +class IOManager; +class IndexFileHandler; +class MemoryPool; +class PkSortedIndexBuilder; +class TableSchema; + +/// Maintains Java-compatible source-backed primary-key indexes for one partition and bucket. +/// +/// This initial implementation builds synchronously during prepare-commit. Synchronous execution +/// deliberately omits Java's scheduling and retry policy while preserving the same source-level +/// reconciliation and payload format. A build failure does not block the data-file commit; readers +/// fall back to scanning the committed data files until a later maintenance attempt succeeds. +class BucketedPrimaryKeyIndexMaintainer { + public: + class Factory { + public: + static Result> Create( + const std::string& root_path, const std::string& branch, + const std::shared_ptr& table_schema, + const std::vector& definitions, + const std::shared_ptr& path_factory, + const std::shared_ptr& index_file_handler, const CoreOptions& options, + const std::shared_ptr& io_manager, bool enable_multi_thread_spill, + const std::shared_ptr& executor, const std::shared_ptr& pool); + + Result> CreateMaintainer( + const BinaryRow& partition, int32_t bucket, + const std::vector>& restored_data_files, + const std::vector>& restored_payloads) const; + + bool Enabled() const { + return !definitions_.empty(); + } + + std::shared_ptr GetIndexFileHandler() const { + return index_file_handler_; + } + + private: + Factory(std::string root_path, std::string branch, + const std::shared_ptr& table_schema, + std::vector definitions, + const std::shared_ptr& path_factory, + const std::shared_ptr& index_file_handler, + const CoreOptions& options, const std::shared_ptr& io_manager, + bool enable_multi_thread_spill, const std::shared_ptr& executor, + const std::shared_ptr& pool) + : root_path_(std::move(root_path)), + branch_(std::move(branch)), + table_schema_(table_schema), + definitions_(std::move(definitions)), + path_factory_(path_factory), + index_file_handler_(index_file_handler), + options_(options), + io_manager_(io_manager), + enable_multi_thread_spill_(enable_multi_thread_spill), + executor_(executor), + pool_(pool) {} + + std::string root_path_; + std::string branch_; + std::shared_ptr table_schema_; + std::vector definitions_; + std::shared_ptr path_factory_; + std::shared_ptr index_file_handler_; + CoreOptions options_; + std::shared_ptr io_manager_; + bool enable_multi_thread_spill_; + std::shared_ptr executor_; + std::shared_ptr pool_; + }; + + Status PrepareCommit(CommitIncrement* increment); + + private: + struct FieldMaintainer { + PrimaryKeyIndexDefinition definition; + std::shared_ptr builder; + }; + + BucketedPrimaryKeyIndexMaintainer( + std::vector fields, + std::map> active_data_files, + std::vector> active_payloads) + : fields_(std::move(fields)), + active_data_files_(std::move(active_data_files)), + active_payloads_(std::move(active_payloads)) {} + + std::vector fields_; + std::map> active_data_files_; + std::vector> active_payloads_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer_test.cpp b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer_test.cpp new file mode 100644 index 000000000..b3f6c18ca --- /dev/null +++ b/src/paimon/core/index/pk/bucketed_primary_key_index_maintainer_test.cpp @@ -0,0 +1,728 @@ +/* + * 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/bucketed_primary_key_index_maintainer.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/catalog/catalog.h" +#include "paimon/catalog/identifier.h" +#include "paimon/commit_context.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/core/global_index/indexed_split_impl.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/pk/primary_key_index_source_policy.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/manifest/file_source.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/table/sink/commit_message_impl.h" +#include "paimon/core/table/source/data_split_impl.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/file_store_commit.h" +#include "paimon/file_store_write.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/read_context.h" +#include "paimon/record_batch.h" +#include "paimon/scan_context.h" +#include "paimon/table/source/plan.h" +#include "paimon/table/source/table_read.h" +#include "paimon/table/source/table_scan.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/write_context.h" + +namespace paimon::test { +namespace { + +constexpr char kCommitUser[] = "pk-index-maintenance-test"; + +std::vector> BTreeIndexFiles(const CommitMessageImpl& message, + bool added) { + std::vector> result; + const std::vector>& data_files = + added ? message.GetNewFilesIncrement().NewIndexFiles() + : message.GetNewFilesIncrement().DeletedIndexFiles(); + const std::vector>& compact_files = + added ? message.GetCompactIncrement().NewIndexFiles() + : message.GetCompactIncrement().DeletedIndexFiles(); + for (const std::shared_ptr& file : data_files) { + if (file != nullptr && file->IndexType() == "btree") { + result.push_back(file); + } + } + for (const std::shared_ptr& file : compact_files) { + if (file != nullptr && file->IndexType() == "btree") { + result.push_back(file); + } + } + return result; +} + +std::vector ExpectedSources( + const std::vector>& files) { + std::vector result; + for (const std::shared_ptr& file : files) { + if (file != nullptr && PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + result.emplace_back(file->file_name, file->row_count); + } + } + std::sort(result.begin(), result.end(), + [](const PrimaryKeyIndexSourceFile& left, const PrimaryKeyIndexSourceFile& right) { + return left.file_name < right.file_name; + }); + return result; +} + +int64_t TotalRows(const std::vector& sources) { + int64_t result = 0; + for (const PrimaryKeyIndexSourceFile& source : sources) { + result += source.row_count; + } + return result; +} + +Result> MakeSourceBackedBTreePayload( + const std::string& file_name, int32_t field_id, const std::shared_ptr& pool) { + constexpr int64_t kRowCount = 1; + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create(/*data_level=*/1, {{"source.data", kRowCount}})); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr source_meta_bytes, source_meta.Serialize(pool)); + return std::make_shared( + "btree", file_name, /*file_size=*/1, kRowCount, /*dv_ranges=*/std::nullopt, + /*external_path=*/std::nullopt, + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/0, field_id, + /*extra_field_ids=*/std::nullopt, /*index_meta=*/nullptr, + source_meta_bytes)); +} + +std::shared_ptr MakeDataEvolutionBTreePayload( + const std::string& file_name, int32_t field_id, const std::shared_ptr& pool) { + constexpr char kSourceMeta[] = + "\x44\x45\x49\x58\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x07"; + auto source_meta = + std::make_shared(std::string(kSourceMeta, sizeof(kSourceMeta) - 1), pool.get()); + return std::make_shared( + "btree", file_name, /*file_size=*/1, /*row_count=*/1, /*dv_ranges=*/std::nullopt, + /*external_path=*/std::nullopt, + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/0, field_id, + /*extra_field_ids=*/std::nullopt, /*index_meta=*/nullptr, source_meta)); +} + +std::shared_ptr MakeMalformedPrimaryKeyBTreePayload( + const std::string& file_name, int32_t field_id, const std::shared_ptr& pool) { + auto source_meta = std::make_shared("malformed", pool.get()); + return std::make_shared( + "btree", file_name, /*file_size=*/1, /*row_count=*/1, /*dv_ranges=*/std::nullopt, + /*external_path=*/std::nullopt, + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/0, field_id, + /*extra_field_ids=*/std::nullopt, /*index_meta=*/nullptr, source_meta)); +} + +Result> ReplaceBTreePayloadSources( + const CommitMessageImpl& message, const std::shared_ptr& payload, + const std::vector& sources, + const std::shared_ptr& pool) { + if (payload == nullptr || !payload->GetGlobalIndexMeta().has_value()) { + return Status::Invalid("Cannot replace source metadata for an invalid BTree payload."); + } + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta original_source_meta, + PrimaryKeyIndexSourceMeta::FromIndexFile(*payload)); + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeyIndexSourceMeta replacement_source_meta, + PrimaryKeyIndexSourceMeta::Create(original_source_meta.DataLevel(), sources)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr source_meta_bytes, + replacement_source_meta.Serialize(pool)); + const GlobalIndexMeta& original_global_meta = payload->GetGlobalIndexMeta().value(); + int64_t row_count = TotalRows(sources); + auto replacement = std::make_shared( + payload->IndexType(), payload->FileName(), payload->FileSize(), row_count, + payload->DvRanges(), payload->ExternalPath(), + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/row_count - 1, + original_global_meta.index_field_id, original_global_meta.extra_field_ids, + original_global_meta.index_meta, source_meta_bytes)); + + std::vector> data_new_indexes = + message.GetNewFilesIncrement().NewIndexFiles(); + std::vector> compact_new_indexes = + message.GetCompactIncrement().NewIndexFiles(); + bool replaced = false; + for (std::vector>* indexes : + {&data_new_indexes, &compact_new_indexes}) { + for (std::shared_ptr& index : *indexes) { + if (index != nullptr && index->FileName() == payload->FileName()) { + index = replacement; + replaced = true; + } + } + } + if (!replaced) { + return Status::Invalid("BTree payload is not present in the commit message."); + } + + const DataIncrement& data_increment = message.GetNewFilesIncrement(); + DataIncrement replacement_data_increment( + std::vector>(data_increment.NewFiles()), + std::vector>(data_increment.DeletedFiles()), + std::vector>(data_increment.ChangelogFiles()), + std::move(data_new_indexes), + std::vector>(data_increment.DeletedIndexFiles())); + const CompactIncrement& compact_increment = message.GetCompactIncrement(); + CompactIncrement replacement_compact_increment( + std::vector>(compact_increment.CompactBefore()), + std::vector>(compact_increment.CompactAfter()), + std::vector>(compact_increment.ChangelogFiles()), + std::move(compact_new_indexes), + std::vector>(compact_increment.DeletedIndexFiles())); + return std::make_shared(message.Partition(), message.Bucket(), + message.TotalBuckets(), replacement_data_increment, + replacement_compact_increment); +} + +} // namespace + +TEST(BucketedPrimaryKeyIndexMaintainerStandaloneTest, + DeletesRestoredBTreePayloadWhenNoDefinitionRemains) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr factory, + BucketedPrimaryKeyIndexMaintainer::Factory::Create( + /*root_path=*/"", /*branch=*/"main", /*table_schema=*/nullptr, + /*definitions=*/{}, /*path_factory=*/nullptr, /*index_file_handler=*/nullptr, options, + /*io_manager=*/nullptr, /*enable_multi_thread_spill=*/false, /*executor=*/nullptr, + GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr stale_payload, + MakeSourceBackedBTreePayload("stale-btree.index", /*field_id=*/7, GetDefaultPool())); + std::shared_ptr data_evolution_payload = MakeDataEvolutionBTreePayload( + "data-evolution-btree.index", /*field_id=*/7, GetDefaultPool()); + std::shared_ptr malformed_pk_payload = MakeMalformedPrimaryKeyBTreePayload( + "malformed-pk-btree.index", /*field_id=*/7, GetDefaultPool()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr maintainer, + factory->CreateMaintainer(BinaryRow::EmptyRow(), /*bucket=*/0, + /*restored_data_files=*/{}, + {stale_payload, data_evolution_payload, malformed_pk_payload})); + CommitIncrement increment(DataIncrement({}, {}, {}), CompactIncrement({}, {}, {}), nullptr); + ASSERT_OK(maintainer->PrepareCommit(&increment)); + ASSERT_EQ(increment.GetNewFilesIncrement().DeletedIndexFiles(), + (std::vector>{stale_payload, malformed_pk_payload})); +} + +TEST(BucketedPrimaryKeyIndexMaintainerStandaloneTest, + PreservesDataEvolutionBTreePayloadForOwnedField) { + constexpr int32_t kFieldId = 7; + PrimaryKeyIndexDefinition definition("value", kFieldId, "btree", + PrimaryKeyIndexDefinition::Family::BTREE, {}); + std::vector fields = { + {std::move(definition), /*builder=*/nullptr}}; + std::shared_ptr data_evolution_payload = + MakeDataEvolutionBTreePayload("data-evolution-btree.index", kFieldId, GetDefaultPool()); + BucketedPrimaryKeyIndexMaintainer maintainer(std::move(fields), /*active_data_files=*/{}, + {data_evolution_payload}); + CommitIncrement increment(DataIncrement({}, {}, {}), CompactIncrement({}, {}, {}), nullptr); + ASSERT_OK(maintainer.PrepareCommit(&increment)); + ASSERT_TRUE(increment.GetNewFilesIncrement().DeletedIndexFiles().empty()); +} + +class BucketedPrimaryKeyIndexMaintainerTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + directory_ = UniqueTestDirectory::Create(); + ASSERT_NE(nullptr, directory_); + schema_ = arrow::schema({arrow::field("id", arrow::int32(), /*nullable=*/false), + arrow::field("value", arrow::utf8(), /*nullable=*/false)}); + options_ = {{Options::BUCKET, "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {Options::FILE_FORMAT, GetParam()}, + {Options::NUM_LEVELS, "3"}, + {Options::PK_BTREE_INDEX_COLUMNS, "value"}, + {Options::TARGET_FILE_ROW_NUM, "2"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + {Options::WRITE_BATCH_SIZE, "2"}}; + if (GetParam() == "orc") { + options_["orc.dictionary-key-size-threshold"] = "1.0"; + options_["orc.read.enable-lazy-decoding"] = "true"; + } + + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, + Catalog::Create(directory_->Str(), {})); + ASSERT_OK(catalog->CreateDatabase("db", {}, /*ignore_if_exists=*/false)); + ASSERT_OK(catalog->CreateTable(Identifier("db", "table"), &c_schema, + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options_, + /*ignore_if_exists=*/false)); + table_path_ = PathUtil::JoinPath(directory_->Str(), "db.db/table"); + } + + Result> MakeBatch(const std::string& json) const { + std::shared_ptr struct_type = arrow::struct_(schema_->fields()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(struct_type, json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetBucket(0).Finish(); + } + + Result> CreateWriter(bool with_temp_directory = true, + bool write_buffer_spillable = true) const { + WriteContextBuilder builder(table_path_, kCommitUser); + std::map writer_options = options_; + writer_options[Options::WRITE_BUFFER_SPILLABLE] = write_buffer_spillable ? "true" : "false"; + builder.SetOptions(writer_options).WithStreamingMode(true); + if (with_temp_directory) { + builder.WithTempDirectory(directory_->Str()); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + return FileStoreWrite::Create(std::move(context)); + } + + Result>> WriteAndPrepare( + const std::string& json, int64_t commit_identifier) const { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, CreateWriter()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch, MakeBatch(json)); + PAIMON_RETURN_NOT_OK(writer->Write(std::move(batch))); + PAIMON_ASSIGN_OR_RAISE(std::vector> messages, + writer->PrepareCommit(/*wait_compaction=*/false, commit_identifier)); + PAIMON_RETURN_NOT_OK(writer->Close()); + return messages; + } + + Result>> CompactAndPrepare( + bool full_compaction, int64_t commit_identifier) const { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, CreateWriter()); + PAIMON_RETURN_NOT_OK(writer->Compact(/*partition=*/{}, /*bucket=*/0, full_compaction)); + PAIMON_ASSIGN_OR_RAISE(std::vector> messages, + writer->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + PAIMON_RETURN_NOT_OK(writer->Close()); + return messages; + } + + Status Commit(const std::vector>& messages, + int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, kCommitUser); + builder.SetOptions(options_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + return commit->Commit(messages, commit_identifier); + } + + std::unique_ptr directory_; + std::shared_ptr schema_; + std::map options_; + std::string table_path_; +}; + +TEST_P(BucketedPrimaryKeyIndexMaintainerTest, BuildsRestoresAndReplacesCompactedLevelPayload) { + ASSERT_OK_AND_ASSIGN(std::vector> initial_messages, + WriteAndPrepare(R"([[4, "b"], [1, "z"], [3, "a"], [2, "y"]])", + /*commit_identifier=*/0)); + ASSERT_EQ(1, initial_messages.size()); + std::shared_ptr initial = + std::dynamic_pointer_cast(initial_messages[0]); + ASSERT_NE(nullptr, initial); + ASSERT_TRUE(BTreeIndexFiles(*initial, /*added=*/true).empty()); + ASSERT_OK(Commit(initial_messages, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(std::vector> first_compact_messages, + CompactAndPrepare(/*full_compaction=*/true, + /*commit_identifier=*/1)); + ASSERT_EQ(1, first_compact_messages.size()); + std::shared_ptr first_compact = + std::dynamic_pointer_cast(first_compact_messages[0]); + ASSERT_NE(nullptr, first_compact); + const std::vector>& first_compact_after = + first_compact->GetCompactIncrement().CompactAfter(); + std::vector first_expected_sources = + ExpectedSources(first_compact_after); + ASSERT_FALSE(first_expected_sources.empty()); + ASSERT_EQ(first_compact_after.size(), first_expected_sources.size()); + + std::vector> first_payloads = + BTreeIndexFiles(*first_compact, /*added=*/true); + ASSERT_EQ(1, first_payloads.size()); + ASSERT_TRUE(BTreeIndexFiles(*first_compact, /*added=*/false).empty()); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta first_source_meta, + PrimaryKeyIndexSourceMeta::FromIndexFile(*first_payloads[0])); + ASSERT_EQ(first_expected_sources, first_source_meta.SourceFiles()); + ASSERT_EQ(first_compact_after[0]->level, first_source_meta.DataLevel()); + int64_t first_row_count = TotalRows(first_expected_sources); + ASSERT_EQ(first_row_count, first_payloads[0]->RowCount()); + ASSERT_TRUE(first_payloads[0]->GetGlobalIndexMeta().has_value()); + ASSERT_EQ(0, first_payloads[0]->GetGlobalIndexMeta()->row_range_start); + ASSERT_EQ(first_row_count - 1, first_payloads[0]->GetGlobalIndexMeta()->row_range_end); + ASSERT_OK(Commit(first_compact_messages, /*commit_identifier=*/1)); + + const std::string indexed_value = "a"; + std::shared_ptr value_predicate = PredicateBuilder::Equal( + /*field_index=*/1, "value", FieldType::STRING, + Literal(FieldType::STRING, indexed_value.data(), indexed_value.size())); + ScanContextBuilder scan_builder(table_path_); + scan_builder.SetPredicate(value_predicate).AddOption(Options::GLOBAL_INDEX_ENABLED, "true"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, table_scan->CreatePlan()); + int32_t indexed_split_count = 0; + int32_t data_split_count = 0; + for (const std::shared_ptr& split : plan->Splits()) { + if (std::dynamic_pointer_cast(split) != nullptr) { + indexed_split_count++; + } else if (std::dynamic_pointer_cast(split) != nullptr) { + data_split_count++; + } + } + ASSERT_GT(indexed_split_count, 0); + ASSERT_EQ(0, data_split_count); + + ReadContextBuilder read_builder(table_path_); + // The source build above keeps lazy decoding enabled. Disable it only for the final data-row + // assertion, whose predicate reader is outside the maintenance path covered by this test. + read_builder.SetReadFieldNames({"id", "value"}) + .SetPredicate(value_predicate) + .AddOption("orc.read.enable-lazy-decoding", "false") + .EnablePredicateFilter(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_NE(nullptr, result); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(result->type(), + {R"([[0, 3, "a"]])"}, &expected) + .ok()); + ASSERT_TRUE(result->Equals(expected)) << result->ToString(); + + ASSERT_OK_AND_ASSIGN(std::vector> unchanged_messages, + CompactAndPrepare(/*full_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_EQ(1, unchanged_messages.size()); + std::shared_ptr unchanged = + std::dynamic_pointer_cast(unchanged_messages[0]); + ASSERT_NE(nullptr, unchanged); + ASSERT_TRUE(unchanged->IsEmpty()); + ASSERT_TRUE(BTreeIndexFiles(*unchanged, /*added=*/true).empty()); + ASSERT_TRUE(BTreeIndexFiles(*unchanged, /*added=*/false).empty()); + + ASSERT_OK_AND_ASSIGN(std::vector> additional_messages, + WriteAndPrepare(R"([[6, "f"], [5, "e"]])", /*commit_identifier=*/2)); + ASSERT_EQ(1, additional_messages.size()); + std::shared_ptr additional = + std::dynamic_pointer_cast(additional_messages[0]); + ASSERT_NE(nullptr, additional); + ASSERT_TRUE(BTreeIndexFiles(*additional, /*added=*/true).empty()); + ASSERT_TRUE(BTreeIndexFiles(*additional, /*added=*/false).empty()); + ASSERT_OK(Commit(additional_messages, /*commit_identifier=*/2)); + + ASSERT_OK_AND_ASSIGN(std::vector> second_compact_messages, + CompactAndPrepare(/*full_compaction=*/true, + /*commit_identifier=*/3)); + ASSERT_EQ(1, second_compact_messages.size()); + std::shared_ptr second_compact = + std::dynamic_pointer_cast(second_compact_messages[0]); + ASSERT_NE(nullptr, second_compact); + std::vector> replacement_payloads = + BTreeIndexFiles(*second_compact, /*added=*/true); + std::vector> deleted_payloads = + BTreeIndexFiles(*second_compact, /*added=*/false); + ASSERT_EQ(1, replacement_payloads.size()); + ASSERT_EQ(1, deleted_payloads.size()); + ASSERT_EQ(first_payloads[0]->FileName(), deleted_payloads[0]->FileName()); + ASSERT_NE(first_payloads[0]->FileName(), replacement_payloads[0]->FileName()); + + std::vector second_expected_sources = + ExpectedSources(second_compact->GetCompactIncrement().CompactAfter()); + ASSERT_FALSE(second_expected_sources.empty()); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta second_source_meta, + PrimaryKeyIndexSourceMeta::FromIndexFile(*replacement_payloads[0])); + ASSERT_EQ(second_expected_sources, second_source_meta.SourceFiles()); + ASSERT_EQ(TotalRows(second_expected_sources), replacement_payloads[0]->RowCount()); +} + +TEST_P(BucketedPrimaryKeyIndexMaintainerTest, RebuildsPartialPayloadAndRetainsItWhenRebuildFails) { + ASSERT_OK_AND_ASSIGN(std::vector> initial_messages, + WriteAndPrepare(R"([[4, "b"], [1, "z"], [3, "a"], [2, "y"]])", + /*commit_identifier=*/0)); + ASSERT_EQ(1, initial_messages.size()); + std::shared_ptr initial = + std::dynamic_pointer_cast(initial_messages[0]); + ASSERT_NE(nullptr, initial); + const DataIncrement& original_data_increment = initial->GetNewFilesIncrement(); + ASSERT_GT(original_data_increment.NewFiles().size(), 1); + std::vector> promoted_sources; + for (const std::shared_ptr& file : original_data_increment.NewFiles()) { + std::shared_ptr promoted = std::make_shared(*file); + promoted->level = 1; + promoted->file_source = FileSource::Compact(); + promoted_sources.push_back(std::move(promoted)); + } + std::vector expected_sources = ExpectedSources(promoted_sources); + ASSERT_GT(expected_sources.size(), 1); + DataIncrement promoted_data_increment( + std::move(promoted_sources), + std::vector>(original_data_increment.DeletedFiles()), + std::vector>(original_data_increment.ChangelogFiles()), + std::vector>(original_data_increment.NewIndexFiles()), + std::vector>(original_data_increment.DeletedIndexFiles())); + std::shared_ptr promoted_message = std::make_shared( + initial->Partition(), initial->Bucket(), initial->TotalBuckets(), promoted_data_increment, + initial->GetCompactIncrement()); + ASSERT_OK(Commit({promoted_message}, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(std::vector> build_messages, + CompactAndPrepare(/*full_compaction=*/false, + /*commit_identifier=*/1)); + ASSERT_EQ(1, build_messages.size()); + std::shared_ptr build = + std::dynamic_pointer_cast(build_messages[0]); + ASSERT_NE(nullptr, build); + ASSERT_TRUE(build->GetCompactIncrement().CompactBefore().empty()); + ASSERT_TRUE(build->GetCompactIncrement().CompactAfter().empty()); + std::vector> payloads = BTreeIndexFiles(*build, /*added=*/true); + ASSERT_EQ(1, payloads.size()); + std::vector partial_sources = {expected_sources[0]}; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr partial_message, + ReplaceBTreePayloadSources(*build, payloads[0], partial_sources, GetDefaultPool())); + ASSERT_OK(Commit({partial_message}, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr failing_writer, + CreateWriter(/*with_temp_directory=*/true, + /*write_buffer_spillable=*/false)); + ASSERT_OK(failing_writer->Compact(/*partition=*/{}, /*bucket=*/0, + /*full_compaction=*/false)); + ASSERT_OK_AND_ASSIGN( + std::vector> failed_rebuild_messages, + failing_writer->PrepareCommit(/*wait_compaction=*/true, /*commit_identifier=*/2)); + ASSERT_OK(failing_writer->Close()); + ASSERT_EQ(1, failed_rebuild_messages.size()); + std::shared_ptr failed_rebuild = + std::dynamic_pointer_cast(failed_rebuild_messages[0]); + ASSERT_NE(nullptr, failed_rebuild); + ASSERT_TRUE(BTreeIndexFiles(*failed_rebuild, /*added=*/true).empty()); + ASSERT_TRUE(BTreeIndexFiles(*failed_rebuild, /*added=*/false).empty()); + + ASSERT_OK_AND_ASSIGN(std::vector> repair_messages, + CompactAndPrepare(/*full_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_EQ(1, repair_messages.size()); + std::shared_ptr repair = + std::dynamic_pointer_cast(repair_messages[0]); + ASSERT_NE(nullptr, repair); + std::vector> replacement_payloads = + BTreeIndexFiles(*repair, /*added=*/true); + std::vector> deleted_payloads = + BTreeIndexFiles(*repair, /*added=*/false); + ASSERT_EQ(1, replacement_payloads.size()); + ASSERT_EQ(1, deleted_payloads.size()); + ASSERT_EQ(payloads[0]->FileName(), deleted_payloads[0]->FileName()); + ASSERT_OK_AND_ASSIGN(PrimaryKeyIndexSourceMeta replacement_source_meta, + PrimaryKeyIndexSourceMeta::FromIndexFile(*replacement_payloads[0])); + ASSERT_EQ(expected_sources, replacement_source_meta.SourceFiles()); +} + +TEST_P(BucketedPrimaryKeyIndexMaintainerTest, ReusesPayloadThatAlsoListsRetiredSources) { + ASSERT_OK_AND_ASSIGN(std::vector> initial_messages, + WriteAndPrepare(R"([[4, "b"], [1, "z"], [3, "a"], [2, "y"]])", + /*commit_identifier=*/0)); + ASSERT_OK(Commit(initial_messages, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(std::vector> compact_messages, + CompactAndPrepare(/*full_compaction=*/true, + /*commit_identifier=*/1)); + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact); + std::vector> payloads = + BTreeIndexFiles(*compact, /*added=*/true); + ASSERT_EQ(1, payloads.size()); + std::vector sources = + ExpectedSources(compact->GetCompactIncrement().CompactAfter()); + ASSERT_FALSE(sources.empty()); + sources.emplace_back("zz-retired.data", /*row_count=*/1); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr retired_source_message, + ReplaceBTreePayloadSources(*compact, payloads[0], sources, GetDefaultPool())); + ASSERT_OK(Commit({retired_source_message}, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> unchanged_messages, + CompactAndPrepare(/*full_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_EQ(1, unchanged_messages.size()); + std::shared_ptr unchanged = + std::dynamic_pointer_cast(unchanged_messages[0]); + ASSERT_NE(nullptr, unchanged); + ASSERT_TRUE(BTreeIndexFiles(*unchanged, /*added=*/true).empty()); + ASSERT_TRUE(BTreeIndexFiles(*unchanged, /*added=*/false).empty()); +} + +TEST_P(BucketedPrimaryKeyIndexMaintainerTest, ContinuesCompactionWhenIndexBuildFails) { + ASSERT_OK_AND_ASSIGN(std::vector> initial_messages, + WriteAndPrepare(R"([[4, "b"], [1, "z"], [3, "a"], [2, "y"]])", + /*commit_identifier=*/0)); + ASSERT_OK(Commit(initial_messages, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateWriter(/*with_temp_directory=*/true, + /*write_buffer_spillable=*/false)); + ASSERT_OK(writer->Compact(/*partition=*/{}, /*bucket=*/0, /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN(std::vector> compact_messages, + writer->PrepareCommit(/*wait_compaction=*/true, + /*commit_identifier=*/1)); + ASSERT_OK(writer->Close()); + + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact); + ASSERT_FALSE(compact->GetCompactIncrement().CompactBefore().empty()); + ASSERT_FALSE(compact->GetCompactIncrement().CompactAfter().empty()); + ASSERT_TRUE(BTreeIndexFiles(*compact, /*added=*/true).empty()); + ASSERT_TRUE(BTreeIndexFiles(*compact, /*added=*/false).empty()); + ASSERT_OK(Commit(compact_messages, /*commit_identifier=*/1)); + + const std::string indexed_value = "a"; + std::shared_ptr value_predicate = PredicateBuilder::Equal( + /*field_index=*/1, "value", FieldType::STRING, + Literal(FieldType::STRING, indexed_value.data(), indexed_value.size())); + ScanContextBuilder scan_builder(table_path_); + scan_builder.SetPredicate(value_predicate).AddOption(Options::GLOBAL_INDEX_ENABLED, "true"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, table_scan->CreatePlan()); + int32_t indexed_split_count = 0; + int32_t data_split_count = 0; + for (const std::shared_ptr& split : plan->Splits()) { + if (std::dynamic_pointer_cast(split) != nullptr) { + indexed_split_count++; + } else if (std::dynamic_pointer_cast(split) != nullptr) { + data_split_count++; + } + } + ASSERT_EQ(0, indexed_split_count); + ASSERT_GT(data_split_count, 0); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetReadFieldNames({"id", "value"}) + .SetPredicate(value_predicate) + .AddOption("orc.read.enable-lazy-decoding", "false") + .EnablePredicateFilter(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(batch_reader.get())); + ASSERT_NE(nullptr, result); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(result->type(), + {R"([[0, 3, "a"]])"}, &expected) + .ok()); + ASSERT_TRUE(result->Equals(expected)) << result->ToString(); + + ASSERT_OK_AND_ASSIGN(std::vector> repair_messages, + CompactAndPrepare(/*full_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_EQ(1, repair_messages.size()); + std::shared_ptr repair = + std::dynamic_pointer_cast(repair_messages[0]); + ASSERT_NE(nullptr, repair); + ASSERT_EQ(1, BTreeIndexFiles(*repair, /*added=*/true).size()); + ASSERT_TRUE(BTreeIndexFiles(*repair, /*added=*/false).empty()); +} + +TEST_P(BucketedPrimaryKeyIndexMaintainerTest, DeletesPayloadAfterLastDefinitionIsRemoved) { + ASSERT_OK_AND_ASSIGN(std::vector> initial_messages, + WriteAndPrepare(R"([[4, "b"], [1, "z"], [3, "a"], [2, "y"]])", + /*commit_identifier=*/0)); + ASSERT_OK(Commit(initial_messages, /*commit_identifier=*/0)); + + ASSERT_OK_AND_ASSIGN(std::vector> compact_messages, + CompactAndPrepare(/*full_compaction=*/true, + /*commit_identifier=*/1)); + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact); + std::vector> payloads = + BTreeIndexFiles(*compact, /*added=*/true); + ASSERT_EQ(1, payloads.size()); + ASSERT_OK(Commit(compact_messages, /*commit_identifier=*/1)); + + std::map evolved_options = options_; + evolved_options.erase(Options::PK_BTREE_INDEX_COLUMNS); + ASSERT_OK_AND_ASSIGN(std::unique_ptr evolved_schema, + TableSchema::Create(/*schema_id=*/1, schema_, /*partition_keys=*/{}, + /*primary_keys=*/{"id"}, evolved_options)); + ASSERT_OK_AND_ASSIGN(std::string schema_json, evolved_schema->ToJsonString()); + ASSERT_OK_AND_ASSIGN(CoreOptions evolved_core_options, CoreOptions::FromMap(evolved_options)); + ASSERT_OK(evolved_core_options.GetFileSystem()->WriteFile( + PathUtil::JoinPath(table_path_, "schema/schema-1"), schema_json, + /*overwrite=*/false)); + + WriteContextBuilder builder(table_path_, kCommitUser); + builder.SetOptions(evolved_options) + .WithStreamingMode(true) + .WithTempDirectory(directory_->Str()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(context))); + ASSERT_OK(writer->Compact(/*partition=*/{}, /*bucket=*/0, /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN(std::vector> cleanup_messages, + writer->PrepareCommit(/*wait_compaction=*/true, + /*commit_identifier=*/2)); + ASSERT_OK(writer->Close()); + ASSERT_EQ(1, cleanup_messages.size()); + std::shared_ptr cleanup = + std::dynamic_pointer_cast(cleanup_messages[0]); + ASSERT_NE(nullptr, cleanup); + std::vector> deleted_payloads = + BTreeIndexFiles(*cleanup, /*added=*/false); + ASSERT_EQ(1, deleted_payloads.size()); + ASSERT_EQ(payloads[0]->FileName(), deleted_payloads[0]->FileName()); + ASSERT_TRUE(BTreeIndexFiles(*cleanup, /*added=*/true).empty()); +} + +INSTANTIATE_TEST_SUITE_P(FileFormats, BucketedPrimaryKeyIndexMaintainerTest, + ::testing::Values("parquet", "orc")); + +} // namespace paimon::test diff --git a/src/paimon/core/index/pk/primary_key_index_definitions.cpp b/src/paimon/core/index/pk/primary_key_index_definitions.cpp index 214338f4b..b63090fbc 100644 --- a/src/paimon/core/index/pk/primary_key_index_definitions.cpp +++ b/src/paimon/core/index/pk/primary_key_index_definitions.cpp @@ -24,6 +24,7 @@ #include #include "fmt/format.h" +#include "paimon/common/global_index/btree/btree_defs.h" #include "paimon/common/utils/object_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/defs.h" @@ -35,7 +36,6 @@ 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"; @@ -180,7 +180,7 @@ Result PrimaryKeyIndexDefinitions::Create(const Tabl PAIMON_ASSIGN_OR_RAISE( IndexOptions definition_options, SortedIndexOptions(options, column, kBTreeOptionFamily, kBTreeAlgorithmPrefix)); - definitions.emplace_back(column, field.Id(), kBTreeIndexType, + definitions.emplace_back(column, field.Id(), BtreeDefs::kIdentifier, PrimaryKeyIndexDefinition::Family::BTREE, std::move(definition_options)); } else if (ObjectUtils::Contains(bitmap_columns, column)) { diff --git a/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.cpp b/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.cpp new file mode 100644 index 000000000..fa8f91620 --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.cpp @@ -0,0 +1,170 @@ +/* + * 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_data_file_reader.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/operation/internal_read_context.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/reader/file_batch_reader.h" + +namespace paimon { + +Result> PkSortedDataFileReader::Create( + const std::string& root_path, const std::shared_ptr& table_schema, + int32_t field_id, const std::shared_ptr& path_factory, + const std::string& branch, const CoreOptions& options, + const std::shared_ptr& executor, const std::shared_ptr& pool) { + std::map read_options = options.ToMap(); + read_options[Options::BRANCH] = branch; + ReadContextBuilder builder(root_path); + builder.SetReadFieldIds({field_id}) + .SetOptions(read_options) + .WithBranch(branch) + .WithFileSystem(options.GetFileSystem()) + .WithExecutor(executor) + .WithMemoryPool(pool) + .EnablePrefetch(false) + .EnablePredicateFilter(false); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, builder.Finish()); + auto shared_read_context = std::shared_ptr(std::move(read_context)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr internal_context, + InternalReadContext::Create(shared_read_context, table_schema, read_options)); + auto shared_internal_context = + std::shared_ptr(std::move(internal_context)); + return std::unique_ptr( + new PkSortedDataFileReader(path_factory, shared_internal_context, pool, executor)); +} + +PkSortedDataFileReader::PkSortedDataFileReader( + const std::shared_ptr& path_factory, + const std::shared_ptr& context, const std::shared_ptr& pool, + const std::shared_ptr& executor) + : RawFileSplitRead(path_factory, context, pool, executor) {} + +Status PkSortedDataFileReader::ReadFile(const BinaryRow& partition, int32_t bucket, + const std::shared_ptr& file, + const BatchConsumer& consumer) const { + if (file == nullptr) { + return Status::Invalid("Primary-key sorted-index source file is null."); + } + if (file->row_count < 0) { + return Status::Invalid(fmt::format("Source file {} has negative row count {}.", + file->file_name, file->row_count)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + path_factory_->CreateDataFilePathFactory(partition, bucket)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + CreateRawFileReaders(partition, {file}, raw_read_schema_, /*predicate=*/nullptr, + /*dv_factory=*/{}, /*row_ranges=*/std::nullopt, data_file_path_factory, + /*extra_format_options=*/{})); + if (readers.size() != 1) { + return Status::Invalid( + fmt::format("Expected one physical reader for source file {}, but got {}.", + file->file_name, readers.size())); + } + std::unique_ptr reader = std::move(readers[0]); + ScopeGuard close_guard([&]() { reader->Close(); }); + PAIMON_ASSIGN_OR_RAISE(uint64_t physical_row_count, reader->GetNumberOfRows()); + if (physical_row_count > static_cast(std::numeric_limits::max()) || + static_cast(physical_row_count) != file->row_count) { + return Status::Invalid(fmt::format( + "Physical row count {} of source file {} does not match metadata row count {}.", + physical_row_count, file->file_name, file->row_count)); + } + + int64_t rows_read = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + break; + } + auto& [batch, bitmap] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (array == nullptr || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + fmt::format("Source file {} did not return a struct batch.", file->file_name)); + } + auto struct_array = checked_pointer_cast(array); + if (struct_array->num_fields() != 1) { + return Status::Invalid( + fmt::format("Source file {} returned {} fields for a single-column index build.", + file->file_name, struct_array->num_fields())); + } + if (static_cast(bitmap.Cardinality()) != struct_array->length()) { + return Status::Invalid( + fmt::format("Source file {} was filtered while building a physical-row index.", + file->file_name)); + } + for (int64_t index = 0; index < struct_array->length(); ++index) { + PAIMON_ASSIGN_OR_RAISE(uint64_t physical_position, + reader->GetPreviousBatchFileRowId(static_cast(index))); + if (physical_position > static_cast(std::numeric_limits::max()) || + static_cast(physical_position) != rows_read + index) { + return Status::Invalid(fmt::format( + "Source file {} returned non-contiguous physical row position {} at row {}.", + file->file_name, physical_position, rows_read + index)); + } + } + PAIMON_RETURN_NOT_OK(consumer(struct_array)); + if (__builtin_add_overflow(rows_read, struct_array->length(), &rows_read)) { + return Status::Invalid("Physical source row count overflows int64."); + } + } + if (rows_read != file->row_count) { + return Status::Invalid( + fmt::format("Read {} physical rows from source file {}, but metadata declares {}.", + rows_read, file->file_name, file->row_count)); + } + return Status::OK(); +} + +Result> PkSortedDataFileReader::ApplyIndexAndDvReaderIfNeeded( + std::unique_ptr&& file_reader, const std::shared_ptr&, + const std::shared_ptr&, const std::shared_ptr& read_schema, + const std::shared_ptr&, DeletionVector::Factory, + const std::optional>&, const std::shared_ptr&) const { + ::ArrowSchema c_read_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); + PAIMON_RETURN_NOT_OK(file_reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + return std::move(file_reader); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.h b/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.h new file mode 100644 index 000000000..de74c2d3b --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_data_file_reader.h @@ -0,0 +1,83 @@ +/* + * 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/type_fwd.h" +#include "paimon/core/deletionvectors/deletion_vector.h" +#include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/result.h" + +namespace arrow { +class StructArray; +} // namespace arrow + +namespace paimon { +class BinaryRow; +class Executor; +class DataFilePathFactory; +class FileBatchReader; +class FileStorePathFactory; +class InternalReadContext; +class MemoryPool; +class Predicate; +class TableSchema; +struct DataFileMeta; + +/// Reads one indexed column from primary-key data files without applying deletion vectors. +/// +/// The reader validates that projected batches contain every physical row in file order before +/// passing them to the callback, so group row ids remain identical to the Java source-backed index +/// contract. +class PkSortedDataFileReader : public RawFileSplitRead { + public: + using BatchConsumer = std::function&)>; + + static Result> Create( + const std::string& root_path, const std::shared_ptr& table_schema, + int32_t field_id, const std::shared_ptr& path_factory, + const std::string& branch, const CoreOptions& options, + const std::shared_ptr& executor, const std::shared_ptr& pool); + + Status ReadFile(const BinaryRow& partition, int32_t bucket, + const std::shared_ptr& file, const BatchConsumer& consumer) const; + + protected: + Result> ApplyIndexAndDvReaderIfNeeded( + std::unique_ptr&& file_reader, const std::shared_ptr& file, + const std::shared_ptr& data_schema, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, + const std::optional>& row_ranges, + const std::shared_ptr& data_file_path_factory) const override; + + private: + PkSortedDataFileReader(const std::shared_ptr& path_factory, + const std::shared_ptr& context, + const std::shared_ptr& pool, + const std::shared_ptr& executor); +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_builder.cpp b/src/paimon/core/index/pksorted/pk_sorted_index_builder.cpp new file mode 100644 index 000000000..8dcbc5b4c --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_builder.cpp @@ -0,0 +1,272 @@ +/* + * 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_builder.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/concatenate.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/casting/casting_utils.h" +#include "paimon/core/global_index/global_index_file_manager.h" +#include "paimon/core/index/pk/primary_key_index_source_file.h" +#include "paimon/core/index/pk/primary_key_index_source_policy.h" +#include "paimon/core/index/pksorted/pk_sorted_data_file_reader.h" +#include "paimon/core/index/pksorted/pk_sorted_index_file.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/mergetree/compact/sort_merge_reader_with_min_heap.h" +#include "paimon/core/mergetree/external_sort_buffer.h" +#include "paimon/core/mergetree/in_memory_sort_buffer.h" +#include "paimon/core/mergetree/sort_buffer.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/fs/file_system.h" +#include "paimon/global_index/io/global_index_file_writer.h" +#include "paimon/record_batch.h" + +namespace paimon { +namespace { + +constexpr char kRowIdFieldName[] = "_PK_INDEX_ROW_ID"; + +class TrackingGlobalIndexFileWriter : public GlobalIndexFileWriter { + public: + explicit TrackingGlobalIndexFileWriter(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + Result NewFileName(const std::string& prefix) const override { + PAIMON_ASSIGN_OR_RAISE(std::string file_name, delegate_->NewFileName(prefix)); + created_file_names_.push_back(file_name); + return file_name; + } + + Result> NewOutputStream( + const std::string& file_name) const override { + return delegate_->NewOutputStream(file_name); + } + + Result GetFileSize(const std::string& file_name) const override { + return delegate_->GetFileSize(file_name); + } + + std::string ToPath(const std::string& file_name) const override { + return delegate_->ToPath(file_name); + } + + void Cleanup(const std::shared_ptr& fs) const { + for (const std::string& file_name : created_file_names_) { + [[maybe_unused]] Status status = fs->Delete(delegate_->ToPath(file_name)); + } + } + + private: + std::shared_ptr delegate_; + mutable std::vector created_file_names_; +}; + +} // namespace + +Result> PkSortedIndexBuilder::Create( + const std::string& root_path, const std::string& branch, const BinaryRow& partition, + int32_t bucket, const std::shared_ptr& table_schema, + const PrimaryKeyIndexDefinition& definition, + const std::shared_ptr& path_factory, const CoreOptions& options, + const std::shared_ptr& io_manager, bool enable_multi_thread_spill, + const std::shared_ptr& executor, const std::shared_ptr& pool) { + if (definition.GetFamily() != PrimaryKeyIndexDefinition::Family::BTREE) { + return Status::Invalid("PkSortedIndexBuilder only supports BTree definitions."); + } + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(definition.FieldId())); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr data_file_reader, + PkSortedDataFileReader::Create(root_path, table_schema, definition.FieldId(), path_factory, + branch, options, executor, pool)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr index_path_factory, + path_factory->CreateIndexFileFactory(partition, bucket)); + return std::unique_ptr(new PkSortedIndexBuilder( + partition, bucket, std::move(field), definition, + std::shared_ptr(std::move(data_file_reader)), + options.GetFileSystem(), std::shared_ptr(std::move(index_path_factory)), + options, io_manager, enable_multi_thread_spill, pool)); +} + +Result> PkSortedIndexBuilder::Build( + const std::vector>& source_files) const { + if (source_files.empty()) { + return Status::Invalid("Cannot build a sorted index for an empty data level."); + } + for (const std::shared_ptr& file : source_files) { + if (file == nullptr) { + return Status::Invalid("A sorted index source file is null."); + } + } + std::vector> ordered_files = source_files; + std::sort( + ordered_files.begin(), ordered_files.end(), + [](const std::shared_ptr& left, const std::shared_ptr& right) { + return left->file_name < right->file_name; + }); + int32_t data_level = ordered_files.front()->level; + std::vector source_metas; + source_metas.reserve(ordered_files.size()); + for (const std::shared_ptr& file : ordered_files) { + if (file == nullptr || file->level != data_level || + !PrimaryKeyIndexSourcePolicy::ShouldRead(*file)) { + return Status::Invalid( + "A sorted index can only cover compacted files from one positive data level."); + } + source_metas.emplace_back(file->file_name, file->row_count); + } + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr unique_comparator, + FieldsComparator::Create({field_}, {0}, + /*is_ascending_order=*/true)); + auto comparator = std::shared_ptr(std::move(unique_comparator)); + DataField row_id_field(std::numeric_limits::max(), + arrow::field(kRowIdFieldName, arrow::int64(), false)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr unique_sequence_comparator, + FieldsComparator::Create({field_, row_id_field}, {1}, /*is_ascending_order=*/true)); + auto sequence_comparator = + std::shared_ptr(std::move(unique_sequence_comparator)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr unique_in_memory_comparator, + FieldsComparator::Create({field_, row_id_field}, {0, 1}, + /*is_ascending_order=*/true)); + auto in_memory_comparator = + std::shared_ptr(std::move(unique_in_memory_comparator)); + auto value_schema = arrow::schema({field_.ArrowField(), row_id_field.ArrowField()}); + // Keep the Arrow memory-pool adapter alive for as long as the sort buffer can retain + // arrays allocated through it. + std::unique_ptr arrow_pool = GetArrowPool(pool_); + auto in_memory_buffer = std::make_unique( + /*last_sequence_number=*/-1, arrow::struct_(value_schema->fields()), + std::vector{field_.Name()}, + /*user_defined_sequence_fields=*/std::vector{kRowIdFieldName}, + /*sequence_fields_ascending=*/true, comparator, options_.GetWriteBufferSize(), pool_, + in_memory_comparator); + std::unique_ptr sort_buffer; + if (options_.GetWriteBufferSpillable() && io_manager_ != nullptr) { + PAIMON_ASSIGN_OR_RAISE( + sort_buffer, + ExternalSortBuffer::Create(std::move(in_memory_buffer), value_schema, {field_.Name()}, + comparator, sequence_comparator, options_, io_manager_, + enable_multi_thread_spill_, pool_)); + } else { + sort_buffer = std::move(in_memory_buffer); + } + ScopeGuard sort_cleanup([&]() { sort_buffer->Clear(); }); + + int64_t rows_buffered = 0; + for (const std::shared_ptr& file : ordered_files) { + Status read_status = data_file_reader_->ReadFile( + partition_, bucket_, file, + [&](const std::shared_ptr& batch) -> Status { + int64_t next_rows_buffered = 0; + if (__builtin_add_overflow(rows_buffered, batch->length(), &next_rows_buffered)) { + return Status::Invalid("Primary-key index row id overflows int64."); + } + std::vector group_ordinals; + group_ordinals.reserve(static_cast(batch->length())); + for (int64_t index = 0; index < batch->length(); ++index) { + group_ordinals.push_back(rows_buffered + index); + } + arrow::Int64Builder row_id_builder(arrow_pool.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_id_builder.AppendValues(group_ordinals)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr row_ids, + row_id_builder.Finish()); + std::shared_ptr indexed_values = batch->field(0); + if (indexed_values->type_id() == arrow::Type::DICTIONARY) { + const auto* dictionary_type = + checked_cast(indexed_values->type().get()); + arrow::Type::type value_type = dictionary_type->value_type()->id(); + if (value_type != arrow::Type::STRING && + value_type != arrow::Type::LARGE_STRING) { + return Status::Invalid(fmt::format( + "Cannot decode dictionary-backed primary-key index field with value " + "type {}.", + dictionary_type->value_type()->ToString())); + } + PAIMON_ASSIGN_OR_RAISE( + indexed_values, + CastingUtils::Cast(indexed_values, field_.ArrowField()->type(), + arrow::compute::CastOptions::Safe(), arrow_pool.get())); + } + // The physical reader owns the pool adapter behind its batch buffers. Copy the + // indexed values into the Build-scoped pool before retaining them in sort_buffer. + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr values, + arrow::Concatenate({indexed_values}, arrow_pool.get())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr sort_batch, + arrow::StructArray::Make({values, row_ids}, {field_.Name(), kRowIdFieldName})); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*sort_batch, &c_array)); + auto record_batch = std::make_unique( + std::map{}, /*bucket=*/0, + std::vector{}, &c_array); + PAIMON_ASSIGN_OR_RAISE(bool has_remaining_quota, + sort_buffer->Write(std::move(record_batch))); + if (!has_remaining_quota) { + return Status::Invalid( + "Primary-key index external-sort quota is exhausted. Configure a " + "temporary directory and sufficient spill capacity."); + } + rows_buffered = next_rows_buffered; + return Status::OK(); + }); + PAIMON_RETURN_NOT_OK(read_status); + } + + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + sort_buffer->CreateReaders()); + auto sorted_reader = std::make_unique( + std::move(readers), comparator, sequence_comparator, + /*merge_function_wrapper=*/nullptr); + auto file_manager = std::make_shared(fs_, index_path_factory_); + auto tracking_writer = std::make_shared(file_manager); + Result> result = PkSortedIndexFile::BuildFromSortedReader( + field_, definition_.IndexType(), definition_.Options(), data_level, source_metas, + std::move(sorted_reader), tracking_writer, index_path_factory_->IsExternalPath(), + options_.GetWriteBatchSize(), pool_); + if (!result.ok()) { + tracking_writer->Cleanup(fs_); + return result.status(); + } + return std::move(result).value(); +} + +Status PkSortedIndexBuilder::DeletePayload(const std::shared_ptr& payload) const { + if (payload == nullptr) { + return Status::OK(); + } + return fs_->Delete(index_path_factory_->ToPath(payload)); +} + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_builder.h b/src/paimon/core/index/pksorted/pk_sorted_index_builder.h new file mode 100644 index 000000000..91706bbb9 --- /dev/null +++ b/src/paimon/core/index/pksorted/pk_sorted_index_builder.h @@ -0,0 +1,96 @@ +/* + * 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/common/data/binary_row.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/result.h" + +namespace paimon { +class Executor; +class FileSystem; +class FileStorePathFactory; +class GlobalIndexFileManager; +class IOManager; +class IndexPathFactory; +class MemoryPool; +class PkSortedDataFileReader; +class TableSchema; +struct DataFileMeta; + +/// Builds one Java-compatible source-backed BTree payload over a complete compacted data level. +class PkSortedIndexBuilder { + public: + static Result> Create( + const std::string& root_path, const std::string& branch, const BinaryRow& partition, + int32_t bucket, const std::shared_ptr& table_schema, + const PrimaryKeyIndexDefinition& definition, + const std::shared_ptr& path_factory, const CoreOptions& options, + const std::shared_ptr& io_manager, bool enable_multi_thread_spill, + const std::shared_ptr& executor, const std::shared_ptr& pool); + + Result> Build( + const std::vector>& source_files) const; + + Status DeletePayload(const std::shared_ptr& payload) const; + + private: + PkSortedIndexBuilder(const BinaryRow& partition, int32_t bucket, DataField field, + PrimaryKeyIndexDefinition definition, + const std::shared_ptr& data_file_reader, + const std::shared_ptr& fs, + const std::shared_ptr& index_path_factory, + const CoreOptions& options, const std::shared_ptr& io_manager, + bool enable_multi_thread_spill, const std::shared_ptr& pool) + : partition_(partition), + bucket_(bucket), + field_(std::move(field)), + definition_(std::move(definition)), + data_file_reader_(data_file_reader), + fs_(fs), + index_path_factory_(index_path_factory), + options_(options), + io_manager_(io_manager), + enable_multi_thread_spill_(enable_multi_thread_spill), + pool_(pool) {} + + BinaryRow partition_; + int32_t bucket_; + DataField field_; + PrimaryKeyIndexDefinition definition_; + std::shared_ptr data_file_reader_; + std::shared_ptr fs_; + std::shared_ptr index_path_factory_; + CoreOptions options_; + std::shared_ptr io_manager_; + bool enable_multi_thread_spill_; + std::shared_ptr pool_; +}; + +} // namespace paimon diff --git a/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp b/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp index 7d3b7aada..6ebae7631 100644 --- a/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp +++ b/src/paimon/core/index/pksorted/pk_sorted_index_file.cpp @@ -20,22 +20,74 @@ #include "paimon/core/index/pksorted/pk_sorted_index_file.h" #include +#include #include #include #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "arrow/record_batch.h" +#include "arrow/type.h" #include "fmt/format.h" +#include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.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/core/io/async_key_value_producer_and_consumer.h" +#include "paimon/core/io/key_value_meta_projection_consumer.h" +#include "paimon/core/io/row_to_arrow_array_converter.h" +#include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/sort_merge_reader.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 { +namespace { + +Result ValidateAndCountSourceRows( + const std::vector& source_files) { + int64_t source_row_count = 0; + for (const PrimaryKeyIndexSourceFile& source_file : source_files) { + if (source_file.row_count < 0 || + __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."); + } + return source_row_count; +} + +Result> FinishIndexFile( + int32_t field_id, const std::string& index_type, int64_t source_row_count, + const PrimaryKeyIndexSourceMeta& source_meta, const std::vector& io_metas, + bool is_external_path, const std::shared_ptr& pool) { + 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) { + PAIMON_ASSIGN_OR_RAISE(Path path, PathUtil::ToPath(io_meta.file_path)); + external_path = path.ToString(); + } + 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 + Result> PkSortedIndexFile::Build( const DataField& field, const std::string& index_type, const std::map& options, int32_t data_level, @@ -47,15 +99,7 @@ Result> PkSortedIndexFile::Build( // external sort buffer and feed the index writer in bounded batches. 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."); - } + PAIMON_ASSIGN_OR_RAISE(int64_t source_row_count, ValidateAndCountSourceRows(source_files)); if (sorted_values == nullptr || sorted_values->length() != source_row_count || static_cast(sorted_ordinals.size()) != source_row_count) { return Status::Invalid( @@ -95,24 +139,96 @@ Result> PkSortedIndexFile::Build( 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())); + return FinishIndexFile(field.Id(), index_type, source_row_count, source_meta, io_metas, + is_external_path, pool); +} + +Result> PkSortedIndexFile::BuildFromSortedReader( + const DataField& field, const std::string& index_type, + const std::map& options, int32_t data_level, + const std::vector& source_files, + std::unique_ptr&& sorted_reader, + const std::shared_ptr& file_writer, bool is_external_path, + int32_t write_batch_size, const std::shared_ptr& pool) { + if (sorted_reader == nullptr) { + return Status::Invalid("Sorted index reader is null."); } - const GlobalIndexIOMeta& io_meta = io_metas[0]; + if (write_batch_size <= 0) { + return Status::Invalid("Sorted index write batch size must be positive."); + } + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create(data_level, source_files)); + PAIMON_ASSIGN_OR_RAISE(int64_t source_row_count, ValidateAndCountSourceRows(source_files)); + 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_schema = arrow::schema({DataField::ConvertDataFieldToArrowField(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(std::shared_ptr source_meta_bytes, source_meta.Serialize(pool)); - std::optional external_path; - if (is_external_path) { - PAIMON_ASSIGN_OR_RAISE(Path path, PathUtil::ToPath(io_meta.file_path)); - external_path = path.ToString(); + auto projection_schema = SpecialFields::CompleteSequenceAndValueKindField(arrow_schema); + auto create_consumer = + [projection_schema, + pool]() -> Result>> { + return KeyValueMetaProjectionConsumer::Create(projection_schema, pool); + }; + std::unique_ptr batch_producer = + std::make_unique(std::move(sorted_reader), write_batch_size); + auto producer = std::make_unique>( + std::move(batch_producer), create_consumer, /*consumer_thread_num=*/1); + ScopeGuard close_guard([&]() { producer->Close(); }); + int64_t rows_written = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, producer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr record_batch, + arrow::ImportRecordBatch(key_value_batch.batch.get(), projection_schema)); + if (record_batch->num_columns() != 3 || + record_batch->column(0)->type_id() != arrow::Type::INT64) { + return Status::Invalid("Sorted index projection produced an invalid batch."); + } + auto sequence_numbers = checked_pointer_cast(record_batch->column(0)); + std::vector ordinals; + ordinals.reserve(static_cast(sequence_numbers->length())); + for (int64_t index = 0; index < sequence_numbers->length(); ++index) { + if (sequence_numbers->IsNull(index)) { + return Status::Invalid("Sorted index row id must not be null."); + } + int64_t ordinal = sequence_numbers->Value(index); + 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)); + } + ordinals.push_back(ordinal); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr values, + arrow::StructArray::Make({record_batch->column(2)}, {field.Name()})); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*values, &c_array)); + ScopeGuard array_guard([&]() { ArrowArrayRelease(&c_array); }); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array, std::move(ordinals))); + if (__builtin_add_overflow(rows_written, values->length(), &rows_written)) { + return Status::Invalid("Sorted index output row count overflows int64."); + } } - 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)); + if (rows_written != source_row_count) { + return Status::Invalid( + fmt::format("Sorted index output row count {} does not match source row count {}.", + rows_written, source_row_count)); + } + PAIMON_ASSIGN_OR_RAISE(std::vector io_metas, writer->Finish()); + return FinishIndexFile(field.Id(), index_type, source_row_count, source_meta, io_metas, + is_external_path, pool); } } // 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 index 1c5089e94..f62fc936b 100644 --- a/src/paimon/core/index/pksorted/pk_sorted_index_file.h +++ b/src/paimon/core/index/pksorted/pk_sorted_index_file.h @@ -34,6 +34,7 @@ #include "paimon/result.h" namespace paimon { +class SortMergeReader; /// Builds one source-backed primary-key index payload for an ordered set of physical data /// files of a single data level. /// @@ -68,6 +69,16 @@ class PkSortedIndexFile { 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); + + /// Builds a payload from a bounded, globally sorted stream. The stream's key contains the + /// indexed field and its sequence number contains the source-group physical row id. + static Result> BuildFromSortedReader( + const DataField& field, const std::string& index_type, + const std::map& options, int32_t data_level, + const std::vector& source_files, + std::unique_ptr&& sorted_reader, + const std::shared_ptr& file_writer, bool is_external_path, + int32_t write_batch_size, const std::shared_ptr& pool); }; } // namespace paimon diff --git a/src/paimon/core/io/key_value_in_memory_record_reader.cpp b/src/paimon/core/io/key_value_in_memory_record_reader.cpp index 9918b328c..158fa5223 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader.cpp +++ b/src/paimon/core/io/key_value_in_memory_record_reader.cpp @@ -18,8 +18,11 @@ #include "paimon/core/io/key_value_in_memory_record_reader.h" +#include #include +#include #include +#include #include "arrow/array/array_base.h" #include "arrow/array/array_nested.h" @@ -59,7 +62,8 @@ KeyValueInMemoryRecordReader::KeyValueInMemoryRecordReader( const std::vector& primary_keys, const std::vector& user_defined_sequence_fields, bool sequence_fields_ascending, const std::shared_ptr& key_comparator, - const std::shared_ptr& pool) + const std::shared_ptr& pool, + const std::shared_ptr& sort_comparator) : last_sequence_num_(last_sequence_num), primary_keys_(primary_keys), user_defined_sequence_fields_(user_defined_sequence_fields), @@ -68,7 +72,8 @@ KeyValueInMemoryRecordReader::KeyValueInMemoryRecordReader( arrow_pool_(GetArrowPool(pool)), value_struct_array_(struct_array), row_kinds_(row_kinds), - key_comparator_(key_comparator) { + key_comparator_(key_comparator), + sort_comparator_(sort_comparator) { assert(value_struct_array_); ArrowUtils::TraverseArray(value_struct_array_); } @@ -110,6 +115,21 @@ void KeyValueInMemoryRecordReader::Close() { Result>> KeyValueInMemoryRecordReader::SortBatch() const { + if (sort_comparator_ != nullptr) { + std::vector indices(static_cast(value_struct_array_->length())); + std::iota(indices.begin(), indices.end(), 0); + std::stable_sort(indices.begin(), indices.end(), [&](uint64_t left, uint64_t right) { + ColumnarRowRef left_row(value_ctx_, left); + ColumnarRowRef right_row(value_ctx_, right); + return sort_comparator_->CompareTo(left_row, right_row) < 0; + }); + arrow::UInt64Builder builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendValues(indices)); + std::shared_ptr sorted_indices; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&sorted_indices)); + return sorted_indices; + } + std::vector sort_keys; sort_keys.reserve(primary_keys_.size() + user_defined_sequence_fields_.size()); for (const auto& name : primary_keys_) { diff --git a/src/paimon/core/io/key_value_in_memory_record_reader.h b/src/paimon/core/io/key_value_in_memory_record_reader.h index 8ffd0c7ba..b6d2a67a7 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader.h +++ b/src/paimon/core/io/key_value_in_memory_record_reader.h @@ -45,14 +45,14 @@ class Metrics; class KeyValueInMemoryRecordReader : public KeyValueRecordReader { public: - KeyValueInMemoryRecordReader(int64_t last_sequence_num, - const std::shared_ptr& struct_array, - const std::vector& row_kinds, - const std::vector& primary_keys, - const std::vector& user_defined_sequence_fields, - bool sequence_fields_ascending, - const std::shared_ptr& key_comparator, - const std::shared_ptr& pool); + KeyValueInMemoryRecordReader( + int64_t last_sequence_num, const std::shared_ptr& struct_array, + const std::vector& row_kinds, + const std::vector& primary_keys, + const std::vector& user_defined_sequence_fields, + bool sequence_fields_ascending, const std::shared_ptr& key_comparator, + const std::shared_ptr& pool, + const std::shared_ptr& sort_comparator = nullptr); class Iterator : public KeyValueRecordReader::Iterator { public: @@ -89,6 +89,7 @@ class KeyValueInMemoryRecordReader : public KeyValueRecordReader { std::shared_ptr value_struct_array_; std::vector row_kinds_; std::shared_ptr key_comparator_; + std::shared_ptr sort_comparator_; std::shared_ptr> sort_indices_; std::shared_ptr key_ctx_; diff --git a/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp b/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp index 7604064ee..01b4149de 100644 --- a/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp +++ b/src/paimon/core/io/key_value_in_memory_record_reader_test.cpp @@ -18,6 +18,7 @@ #include "paimon/core/io/key_value_in_memory_record_reader.h" +#include #include #include #include @@ -254,6 +255,49 @@ TEST_F(KeyValueInMemoryRecordReaderTest, TestUserDefinedSequenceFieldsDescending KeyValueChecker::CheckResult(expected, results, /*key_arity=*/2, /*value_arity=*/5); } +TEST_F(KeyValueInMemoryRecordReaderTest, TestComparatorBackedFloatingPointSort) { + std::vector fields = { + DataField(0, arrow::field("key", arrow::float32(), /*nullable=*/false)), + DataField(1, arrow::field("row_id", arrow::int64(), /*nullable=*/false))}; + arrow::FloatBuilder key_builder; + ASSERT_TRUE(key_builder + .AppendValues({std::numeric_limits::quiet_NaN(), 0.0F, 1.0F, -0.0F, + -std::numeric_limits::infinity(), + std::numeric_limits::infinity()}) + .ok()); + arrow::Int64Builder row_id_builder; + ASSERT_TRUE(row_id_builder.AppendValues({0, 1, 2, 3, 4, 5}).ok()); + std::shared_ptr keys = key_builder.Finish().ValueOrDie(); + std::shared_ptr row_ids = row_id_builder.Finish().ValueOrDie(); + std::shared_ptr source = + arrow::StructArray::Make({keys, row_ids}, {fields[0].ArrowField(), fields[1].ArrowField()}) + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({fields[0]}, {0}, + /*is_ascending_order=*/true)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr sort_comparator, + FieldsComparator::Create(fields, {0, 1}, + /*is_ascending_order=*/true)); + auto reader = std::make_unique( + /*last_sequence_num=*/0, source, std::vector(), + std::vector{"key"}, std::vector{"row_id"}, + /*sequence_fields_ascending=*/true, key_comparator, pool_, sort_comparator); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr iterator, + reader->NextBatch()); + std::vector sequence_numbers; + while (true) { + ASSERT_OK_AND_ASSIGN(bool has_next, iterator->HasNext()); + if (!has_next) { + break; + } + ASSERT_OK_AND_ASSIGN(KeyValue key_value, iterator->Next()); + sequence_numbers.push_back(key_value.sequence_number); + } + ASSERT_EQ((std::vector{4, 3, 1, 2, 5, 0}), sequence_numbers); +} + TEST_F(KeyValueInMemoryRecordReaderTest, TestNonExistPK) { std::vector fields = {DataField(0, arrow::field("k0", arrow::int32())), DataField(1, arrow::field("k1", arrow::int32())), diff --git a/src/paimon/core/mergetree/in_memory_sort_buffer.cpp b/src/paimon/core/mergetree/in_memory_sort_buffer.cpp index 5f53d5984..493038b78 100644 --- a/src/paimon/core/mergetree/in_memory_sort_buffer.cpp +++ b/src/paimon/core/mergetree/in_memory_sort_buffer.cpp @@ -42,13 +42,15 @@ InMemorySortBuffer::InMemorySortBuffer(int64_t last_sequence_number, bool sequence_fields_ascending, const std::shared_ptr& key_comparator, uint64_t write_buffer_size, - const std::shared_ptr& pool) + const std::shared_ptr& pool, + const std::shared_ptr& sort_comparator) : pool_(pool), value_type_(value_type), trimmed_primary_keys_(trimmed_primary_keys), user_defined_sequence_fields_(user_defined_sequence_fields), sequence_fields_ascending_(sequence_fields_ascending), key_comparator_(key_comparator), + sort_comparator_(sort_comparator), write_buffer_size_(write_buffer_size), next_sequence_number_(last_sequence_number + 1) {} @@ -104,7 +106,7 @@ Result>> InMemorySortBuffer::C auto in_memory_reader = std::make_unique( buffered_batch.first_sequence_number, buffered_batch.struct_array, buffered_batch.row_kinds, trimmed_primary_keys_, user_defined_sequence_fields_, - sequence_fields_ascending_, key_comparator_, pool_); + sequence_fields_ascending_, key_comparator_, pool_, sort_comparator_); readers.push_back(std::move(in_memory_reader)); } return readers; diff --git a/src/paimon/core/mergetree/in_memory_sort_buffer.h b/src/paimon/core/mergetree/in_memory_sort_buffer.h index 3632bb86e..92e3ef5a2 100644 --- a/src/paimon/core/mergetree/in_memory_sort_buffer.h +++ b/src/paimon/core/mergetree/in_memory_sort_buffer.h @@ -57,7 +57,8 @@ class InMemorySortBuffer : public SortBuffer { const std::vector& user_defined_sequence_fields, bool sequence_fields_ascending, const std::shared_ptr& key_comparator, - uint64_t write_buffer_size, const std::shared_ptr& pool); + uint64_t write_buffer_size, const std::shared_ptr& pool, + const std::shared_ptr& sort_comparator = nullptr); void Clear() override; uint64_t GetMemorySize() const override; @@ -78,6 +79,7 @@ class InMemorySortBuffer : public SortBuffer { const std::vector user_defined_sequence_fields_; const bool sequence_fields_ascending_; const std::shared_ptr key_comparator_; + const std::shared_ptr sort_comparator_; const uint64_t write_buffer_size_; std::vector buffered_batches_; diff --git a/src/paimon/core/operation/abstract_file_store_write.cpp b/src/paimon/core/operation/abstract_file_store_write.cpp index d161e5d5f..6d827ffa5 100644 --- a/src/paimon/core/operation/abstract_file_store_write.cpp +++ b/src/paimon/core/operation/abstract_file_store_write.cpp @@ -60,6 +60,8 @@ AbstractFileStoreWrite::AbstractFileStoreWrite( const std::shared_ptr& write_schema, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, + const std::shared_ptr& + primary_key_index_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, const std::shared_ptr& executor, const std::shared_ptr& pool) @@ -75,6 +77,7 @@ AbstractFileStoreWrite::AbstractFileStoreWrite( table_schema_(table_schema), partition_schema_(partition_schema), dv_maintainer_factory_(dv_maintainer_factory), + primary_key_index_maintainer_factory_(primary_key_index_maintainer_factory), io_manager_(io_manager), options_(options), compact_executor_(CreateDefaultExecutor()), @@ -208,6 +211,29 @@ Result>> AbstractFileStoreWrite::Prep compact_increment.AddNewIndexFiles({dv_index_file_meta.value()}); } } + if (writer_container.primary_key_index_maintainer) { + Status index_status = + writer_container.primary_key_index_maintainer->PrepareCommit(&increment); + if (!index_status.ok()) { + if (compact_deletion_file) { + const auto& new_index_files = + increment.GetCompactIncrement().NewIndexFiles(); + for (const std::shared_ptr& index_file : new_index_files) { + if (index_file != nullptr && + index_file->IndexType() == + DeletionVectorsIndexFile::DELETION_VECTORS_INDEX) { + PAIMON_ASSIGN_OR_RAISE( + std::string index_path, + dv_maintainer_factory_->GetIndexFileHandler()->FilePath( + partition, bucket, index_file)); + [[maybe_unused]] Status cleanup_status = + options_.GetFileSystem()->Delete(index_path); + } + } + } + return index_status; + } + } auto committable = std::make_shared( partition, bucket, writer_container.total_buckets, increment.GetNewFilesIncrement(), @@ -375,6 +401,8 @@ Result> AbstractFileStoreWrite::ScanExistingFileMe std::shared_ptr index_file_handler; if (dv_maintainer_factory_) { index_file_handler = dv_maintainer_factory_->GetIndexFileHandler(); + } else if (primary_key_index_maintainer_factory_) { + index_file_handler = primary_key_index_maintainer_factory_->GetIndexFileHandler(); } // Paimon Java currently drops value stats during writer restore. This is a known bug: a // restored file can become a compact-after ADD via metadata-only level upgrade and lose its @@ -383,7 +411,8 @@ Result> AbstractFileStoreWrite::ScanExistingFileMe FileSystemWriteRestore restore(snapshot_manager_, std::move(scan), index_file_handler); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr restore_files, - restore.GetRestoreFiles(partition, bucket, dv_maintainer_factory_ != nullptr)); + restore.GetRestoreFiles(partition, bucket, dv_maintainer_factory_ != nullptr, + primary_key_index_maintainer_factory_ != nullptr)); std::optional restored_total_buckets = restore_files->TotalBuckets(); int32_t total_buckets = GetDefaultBucketNum(); @@ -427,17 +456,28 @@ Result> AbstractFileStoreWrite::GetWriter(const Bin dv_maintainer_factory_->Create(partition, bucket, restored->DeleteVectorsIndex())); } + std::shared_ptr primary_key_index_maintainer; + if (primary_key_index_maintainer_factory_) { + PAIMON_ASSIGN_OR_RAISE( + primary_key_index_maintainer, + primary_key_index_maintainer_factory_->CreateMaintainer( + partition, bucket, restore_data_files, restored->PrimaryKeyIndexPayloads())); + } + PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, CreateWriter(partition, bucket, restore_data_files, max_sequence_number, dv_maintainer)); int32_t total_buckets = restored->TotalBuckets().value_or(GetDefaultBucketNum()); if (partition_iter == writers_.end()) { - writers_.emplace(partition, - std::unordered_map>( - {{bucket, WriterContainer(writer, total_buckets)}})); + writers_.emplace( + partition, std::unordered_map>( + {{bucket, WriterContainer(writer, total_buckets, + primary_key_index_maintainer)}})); } else { - partition_iter->second.emplace(bucket, WriterContainer(writer, total_buckets)); + partition_iter->second.emplace( + bucket, + WriterContainer(writer, total_buckets, primary_key_index_maintainer)); } writer_memory_manager_->RegisterWriter(writer.get()); diff --git a/src/paimon/core/operation/abstract_file_store_write.h b/src/paimon/core/operation/abstract_file_store_write.h index f43046bfd..0a0a9d340 100644 --- a/src/paimon/core/operation/abstract_file_store_write.h +++ b/src/paimon/core/operation/abstract_file_store_write.h @@ -32,6 +32,7 @@ #include "paimon/common/io/cache/cache_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" +#include "paimon/core/index/pk/bucketed_primary_key_index_maintainer.h" #include "paimon/core/memory/writer_memory_manager.h" #include "paimon/file_store_write.h" #include "paimon/logging.h" @@ -79,6 +80,8 @@ class AbstractFileStoreWrite : public FileStoreWrite { const std::shared_ptr& write_schema, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, + const std::shared_ptr& + primary_key_index_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, const std::shared_ptr& executor, const std::shared_ptr& pool); @@ -102,11 +105,16 @@ class AbstractFileStoreWrite : public FileStoreWrite { struct WriterContainer { public: WriterContainer() = default; - WriterContainer(const std::shared_ptr& writer, int32_t total_buckets) - : writer(writer), total_buckets(total_buckets) {} + WriterContainer( + const std::shared_ptr& writer, int32_t total_buckets, + const std::shared_ptr& primary_key_index_maintainer) + : writer(writer), + total_buckets(total_buckets), + primary_key_index_maintainer(primary_key_index_maintainer) {} std::shared_ptr writer; int64_t last_modified_commit_identifier = std::numeric_limits::min(); int32_t total_buckets = -1; + std::shared_ptr primary_key_index_maintainer; }; protected: @@ -139,6 +147,8 @@ class AbstractFileStoreWrite : public FileStoreWrite { std::shared_ptr table_schema_; std::shared_ptr partition_schema_; std::shared_ptr dv_maintainer_factory_; + std::shared_ptr + primary_key_index_maintainer_factory_; std::shared_ptr io_manager_; std::shared_ptr cache_manager_; std::unique_ptr writer_memory_manager_; diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index f660093a2..f113133a6 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -81,10 +81,11 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool) - : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, - root_path, table_schema, schema, write_schema, partition_schema, - dv_maintainer_factory, io_manager, options, ignore_previous_files, - is_streaming_mode, ignore_num_bucket_check, executor, pool), + : AbstractFileStoreWrite( + file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, + table_schema, schema, write_schema, partition_schema, dv_maintainer_factory, + /*primary_key_index_maintainer_factory=*/nullptr, io_manager, options, + ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, executor, pool), realtime_context_(realtime_context), logger_(Logger::GetLogger("AppendOnlyFileStoreWrite")) { write_cols_ = write_schema->field_names(); diff --git a/src/paimon/core/operation/expire_snapshots.cpp b/src/paimon/core/operation/expire_snapshots.cpp index a5bd25053..2de2eb38b 100644 --- a/src/paimon/core/operation/expire_snapshots.cpp +++ b/src/paimon/core/operation/expire_snapshots.cpp @@ -19,12 +19,12 @@ #include "paimon/core/operation/expire_snapshots.h" #include -#include #include #include #include #include #include +#include #include #include "fmt/format.h" @@ -33,15 +33,21 @@ #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" +#include "paimon/core/index/index_path_factory.h" +#include "paimon/core/manifest/file_entry.h" #include "paimon/core/manifest/file_kind.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_entry.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/snapshot.h" +#include "paimon/core/utils/branch_manager.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/core/utils/tag_manager.h" #include "paimon/fs/file_system.h" namespace paimon { @@ -50,12 +56,14 @@ ExpireSnapshots::ExpireSnapshots(const std::shared_ptr& snapsho const std::shared_ptr& path_factory, const std::shared_ptr& manifest_list, const std::shared_ptr& manifest_file, + const std::shared_ptr& index_manifest_file, const std::shared_ptr& fs, const ExpireConfig& config, bool realtime_enabled, const std::shared_ptr& executor) : snapshot_manager_(snapshot_manager), path_factory_(path_factory), manifest_list_(manifest_list), manifest_file_(manifest_file), + index_manifest_file_(index_manifest_file), fs_(fs), config_(config), realtime_enabled_(realtime_enabled), @@ -123,6 +131,13 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, // TODO(jinli.zjw): write earliest hint return 0; } + PAIMON_ASSIGN_OR_RAISE(std::vector branches, + BranchManager::ListBranches(fs_, snapshot_manager_->RootPath())); + if (branches.size() > 1) { + return Status::NotImplemented( + "Snapshot expiration is disabled while another branch exists because cross-branch " + "file retention is not supported."); + } int64_t begin_inclusive_id = earliest_snapshot_id; for (int64_t id = end_exclusive_id - 1; id >= earliest_snapshot_id; id--) { PAIMON_ASSIGN_OR_RAISE(bool exist, snapshot_manager_->SnapshotExists(id)); @@ -134,6 +149,12 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, PAIMON_LOG_DEBUG(logger_, "Snapshot expire range is [%ld, %ld]", begin_inclusive_id, end_exclusive_id); + PAIMON_ASSIGN_OR_RAISE(std::vector tagged_snapshots, GetTaggedSnapshots()); + auto next_tag = tagged_snapshots.begin(); + const Snapshot* previous_tag = nullptr; + std::optional> tagged_data_files; + const std::set no_retained_data_files; + // Since the data file deletion information for each snapshot is recorded in the delta part of // the next snapshot, it is necessary to check the next snapshot. Otherwise, its data files will // not be deleted in this round. @@ -144,7 +165,35 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, continue; } PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(id)); - PAIMON_RETURN_NOT_OK(CleanUnusedDataFiles(snapshot.DeltaManifestList())); + bool tag_changed = false; + while (next_tag != tagged_snapshots.end() && next_tag->Id() < id) { + previous_tag = &*next_tag; + ++next_tag; + tag_changed = true; + } + if (tag_changed) { + Result> tagged_data_files_result = + GetTaggedDataFiles(*previous_tag); + if (!tagged_data_files_result.ok()) { + PAIMON_LOG_WARN(logger_, + "Skip cleaning data files of snapshot #%ld because the data files " + "referenced by tag snapshot #%ld could not be loaded. Snapshot " + "metadata expiration will continue, so skipped data files may " + "remain for orphan cleanup. %s", + id, previous_tag->Id(), + tagged_data_files_result.status().ToString().c_str()); + tagged_data_files.reset(); + } else { + tagged_data_files = std::move(tagged_data_files_result).value(); + } + } + if (previous_tag != nullptr && !tagged_data_files) { + continue; + } + const std::set& retained_data_files = + tagged_data_files ? tagged_data_files.value() : no_retained_data_files; + PAIMON_RETURN_NOT_OK( + CleanUnusedDataFiles(snapshot.DeltaManifestList(), retained_data_files)); } // TODO(jinli.zjw): support delete changelog files @@ -159,6 +208,10 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, std::vector retained_snapshots; PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(end_exclusive_id)); retained_snapshots.push_back(snapshot); + std::vector retained_tag_snapshots = + GetTagSnapshotsToRetain(tagged_snapshots, begin_inclusive_id, end_exclusive_id); + retained_snapshots.insert(retained_snapshots.end(), retained_tag_snapshots.begin(), + retained_tag_snapshots.end()); std::set retained_offset_files; if (realtime_enabled_) { PAIMON_ASSIGN_OR_RAISE(std::vector all_snapshots, @@ -185,6 +238,7 @@ Result ExpireSnapshots::ExpireUntil(int64_t earliest_snapshot_id, continue; } PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(id)); + PAIMON_RETURN_NOT_OK(CleanUnusedIndexManifest(snapshot.IndexManifest(), &skipping_sets)); PAIMON_RETURN_NOT_OK(CleanUnusedManifests(snapshot.BaseManifestList(), skipping_sets)); PAIMON_RETURN_NOT_OK(CleanUnusedManifests(snapshot.DeltaManifestList(), skipping_sets)); if (realtime_enabled_) { @@ -290,7 +344,58 @@ Status ExpireSnapshots::CleanUnusedManifests(const std::string& manifest_list_na return Status::OK(); } -Status ExpireSnapshots::CleanUnusedDataFiles(const std::string& manifest_list_name) { +Status ExpireSnapshots::CleanUnusedIndexManifest(const std::optional& index_manifest, + std::set* skipping_manifest_set) { + if (!index_manifest || index_manifest->empty() || + skipping_manifest_set->count(index_manifest.value()) > 0) { + return Status::OK(); + } + if (index_manifest_file_ == nullptr) { + return Status::Invalid("index manifest file is null"); + } + + std::vector entries; + Status read_status = + index_manifest_file_->ReadIfFileExist(index_manifest.value(), /*filter=*/nullptr, &entries); + if (read_status.IsNotExist()) { + return Status::OK(); + } + PAIMON_RETURN_NOT_OK(read_status); + + std::vector> index_files_to_delete; + std::set planned_file_names; + for (const IndexManifestEntry& entry : entries) { + const std::string& file_name = entry.index_file->FileName(); + if (skipping_manifest_set->count(file_name) > 0 || + !planned_file_names.insert(file_name).second) { + continue; + } + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr index_path_factory, + path_factory_->CreateIndexFileFactory(entry.partition, entry.bucket)); + index_files_to_delete.emplace_back(file_name, index_path_factory->ToPath(entry.index_file), + entry.index_file->ExternalPath().has_value()); + } + + for (const auto& [file_name, file_path, is_external_path] : index_files_to_delete) { + Status delete_status = fs_->Delete(file_path); + if (!delete_status.ok() && is_external_path) { + PAIMON_ASSIGN_OR_RAISE(bool exists, fs_->Exists(file_path)); + if (exists) { + return delete_status.WithMessage("Failed to delete external index payload '", + file_path, "': ", delete_status.message()); + } + } + skipping_manifest_set->insert(file_name); + } + + skipping_manifest_set->insert(index_manifest.value()); + index_manifest_file_->DeleteQuietly(index_manifest.value()); + return Status::OK(); +} + +Status ExpireSnapshots::CleanUnusedDataFiles(const std::string& manifest_list_name, + const std::set& retained_data_files) { std::vector manifest_file_metas; auto status = manifest_list_->Read(manifest_list_name, nullptr, &manifest_file_metas); if (status.ok()) { @@ -308,6 +413,9 @@ Status ExpireSnapshots::CleanUnusedDataFiles(const std::string& manifest_list_na PAIMON_RETURN_NOT_OK(GetDataFilesToDelete(manifest_entries, &data_files_to_delete)); } } + for (const std::string& retained_data_file : retained_data_files) { + data_files_to_delete.erase(retained_data_file); + } std::vector> futures; ScopeGuard guard([&futures]() { Wait(futures); }); @@ -354,11 +462,8 @@ Status ExpireSnapshots::GetManifestSkippingSet(const std::vector& reta for (const auto& manifest : manifests) { skipping_manifest_set->insert(manifest.FileName()); } - // TODO(jinli.zjw): skip index manifests - if (snapshot.IndexManifest() && snapshot.IndexManifest().value() != "") { - assert(false); - return Status::NotImplemented("do not support expire snapshot with index manifest"); - } + PAIMON_RETURN_NOT_OK( + AddIndexManifestToSkippingSet(snapshot.IndexManifest(), skipping_manifest_set)); if (snapshot.Statistics()) { skipping_manifest_set->insert(snapshot.Statistics().value()); } @@ -366,4 +471,85 @@ Status ExpireSnapshots::GetManifestSkippingSet(const std::vector& reta return Status::OK(); } +Result> ExpireSnapshots::GetTaggedSnapshots() const { + TagManager tag_manager(fs_, snapshot_manager_->RootPath(), snapshot_manager_->Branch()); + PAIMON_ASSIGN_OR_RAISE(std::vector tag_names, tag_manager.ListTagNames()); + + std::vector tagged_snapshots; + tagged_snapshots.reserve(tag_names.size()); + for (const std::string& tag_name : tag_names) { + PAIMON_ASSIGN_OR_RAISE(std::optional tag, tag_manager.Get(tag_name)); + if (!tag) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(Snapshot tagged_snapshot, tag->TrimToSnapshot()); + tagged_snapshots.push_back(std::move(tagged_snapshot)); + } + std::sort(tagged_snapshots.begin(), tagged_snapshots.end(), + [](const Snapshot& lhs, const Snapshot& rhs) { return lhs.Id() < rhs.Id(); }); + return tagged_snapshots; +} + +std::vector ExpireSnapshots::GetTagSnapshotsToRetain( + const std::vector& tagged_snapshots, int64_t begin_inclusive_id, + int64_t end_exclusive_id) const { + auto right = std::lower_bound( + tagged_snapshots.begin(), tagged_snapshots.end(), end_exclusive_id, + [](const Snapshot& snapshot, int64_t snapshot_id) { return snapshot.Id() < snapshot_id; }); + if (right == tagged_snapshots.begin()) { + return std::vector(); + } + auto left = std::upper_bound( + tagged_snapshots.begin(), right, begin_inclusive_id, + [](int64_t snapshot_id, const Snapshot& snapshot) { return snapshot_id < snapshot.Id(); }); + if (left != tagged_snapshots.begin()) { + --left; + } + return std::vector(left, right); +} + +Result> ExpireSnapshots::GetTaggedDataFiles( + const Snapshot& tagged_snapshot) const { + std::vector manifests; + PAIMON_RETURN_NOT_OK(manifest_list_->ReadDataManifests(tagged_snapshot, &manifests)); + + std::vector unmerged_entries; + for (const ManifestFileMeta& manifest : manifests) { + std::vector entries; + PAIMON_RETURN_NOT_OK( + manifest_file_->Read(manifest.FileName(), /*filter=*/nullptr, &entries)); + unmerged_entries.insert(unmerged_entries.end(), entries.begin(), entries.end()); + } + std::vector merged_entries; + PAIMON_RETURN_NOT_OK(FileEntry::MergeEntries(unmerged_entries, &merged_entries)); + + std::set tagged_data_files; + for (const ManifestEntry& entry : merged_entries) { + PAIMON_ASSIGN_OR_RAISE(std::string bucket_path, + path_factory_->BucketPath(entry.Partition(), entry.Bucket())); + tagged_data_files.insert(PathUtil::JoinPath(bucket_path, entry.FileName())); + } + return tagged_data_files; +} + +Status ExpireSnapshots::AddIndexManifestToSkippingSet( + const std::optional& index_manifest, + std::set* skipping_manifest_set) const { + if (!index_manifest || index_manifest->empty()) { + return Status::OK(); + } + if (index_manifest_file_ == nullptr) { + return Status::Invalid("index manifest file is null"); + } + + skipping_manifest_set->insert(index_manifest.value()); + std::vector entries; + PAIMON_RETURN_NOT_OK( + index_manifest_file_->Read(index_manifest.value(), /*filter=*/nullptr, &entries)); + for (const IndexManifestEntry& entry : entries) { + skipping_manifest_set->insert(entry.index_file->FileName()); + } + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/core/operation/expire_snapshots.h b/src/paimon/core/operation/expire_snapshots.h index 238f599ed..a21519b27 100644 --- a/src/paimon/core/operation/expire_snapshots.h +++ b/src/paimon/core/operation/expire_snapshots.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -38,6 +39,7 @@ class Snapshot; class SnapshotManager; class FileStorePathFactory; class FileSystem; +class IndexManifestFile; class ManifestEntry; class ManifestList; class ManifestFile; @@ -50,6 +52,7 @@ class ExpireSnapshots { const std::shared_ptr& path_factory, const std::shared_ptr& manifest_list, const std::shared_ptr& manifest_file, + const std::shared_ptr& index_manifest_file, const std::shared_ptr& fs, const ExpireConfig& config, bool realtime_enabled, const std::shared_ptr& executor); @@ -58,20 +61,31 @@ class ExpireSnapshots { private: Result ExpireUntil(int64_t earliest_snapshot_id, int64_t end_exclusive_id); - Status CleanUnusedDataFiles(const std::string& manifest_list_name); + Status CleanUnusedDataFiles(const std::string& manifest_list_name, + const std::set& retained_data_files); Status CleanUnusedManifests(const std::string& manifest_list_name, const std::set& skipping_sets); + Status CleanUnusedIndexManifest(const std::optional& index_manifest, + std::set* skipping_manifest_set); Status CleanEmptyDirectories(); Status GetDataFilesToDelete(const std::vector& data_file_entries, std::map* data_files_to_delete) const; Status GetManifestSkippingSet(const std::vector& retained_snapshots, std::set* skipping_manifest_set) const; + Result> GetTaggedSnapshots() const; + std::vector GetTagSnapshotsToRetain(const std::vector& tagged_snapshots, + int64_t begin_inclusive_id, + int64_t end_exclusive_id) const; + Result> GetTaggedDataFiles(const Snapshot& tagged_snapshot) const; + Status AddIndexManifestToSkippingSet(const std::optional& index_manifest, + std::set* skipping_manifest_set) const; bool TryDeleteEmptyDirectory(const std::string& path) const; std::shared_ptr snapshot_manager_; std::shared_ptr path_factory_; std::shared_ptr manifest_list_; std::shared_ptr manifest_file_; + std::shared_ptr index_manifest_file_; std::shared_ptr fs_; ExpireConfig config_; bool realtime_enabled_; diff --git a/src/paimon/core/operation/expire_snapshots_test.cpp b/src/paimon/core/operation/expire_snapshots_test.cpp index ffdd30feb..2d5bde9c0 100644 --- a/src/paimon/core/operation/expire_snapshots_test.cpp +++ b/src/paimon/core/operation/expire_snapshots_test.cpp @@ -20,22 +20,33 @@ #include #include +#include #include #include "arrow/type.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/data/binary_row_writer.h" +#include "paimon/common/utils/path_util.h" #include "paimon/core/core_options.h" +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/index_path_factory.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_kind.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_entry.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/stats/simple_stats.h" +#include "paimon/core/tag/tag.h" #include "paimon/core/utils/field_mapping.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/core/utils/tag_manager.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" #include "paimon/executor.h" @@ -45,6 +56,25 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +class DeleteFailingFileSystem : public LocalFileSystem { + public: + explicit DeleteFailingFileSystem(std::string failed_path) + : failed_path_(std::move(failed_path)) {} + + Status Delete(const std::string& path, bool recursive = true) const override { + if (path == failed_path_) { + return Status::IOError("injected delete failure"); + } + return LocalFileSystem::Delete(path, recursive); + } + + private: + std::string failed_path_; +}; + +} // namespace class ExpireSnapshotsTest : public testing::Test { public: @@ -70,9 +100,11 @@ class ExpireSnapshotsTest : public testing::Test { schema_ = arrow::schema(fields); ASSERT_OK_AND_ASSIGN(partition_schema_, FieldMapping::GetPartitionSchema(schema_, partition_keys_)); - fs_ = std::make_shared(); + dir_ = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir_); + fs_ = dir_->GetFileSystem(); - test_data_path_ = "tmp"; + test_data_path_ = dir_->Str(); path_factory_ = CreateFactory(test_data_path_); ASSERT_OK_AND_ASSIGN( @@ -85,6 +117,11 @@ class ExpireSnapshotsTest : public testing::Test { ManifestFile::Create(fs_, options.GetManifestFormat(), options.GetManifestCompression(), path_factory_, options.GetManifestTargetFileSize(), mem_pool_, options, partition_schema_)); + + ASSERT_OK_AND_ASSIGN(index_manifest_file_, + IndexManifestFile::Create( + fs_, options.GetManifestFormat(), options.GetManifestCompression(), + path_factory_, options.GetBucket(), mem_pool_, options)); } void TearDown() override {} @@ -152,7 +189,36 @@ class ExpireSnapshotsTest : public testing::Test { return ManifestEntry(kind, row, bucket, /*total_buckets=*/3, data_file_meta); } + Result CreateSourceBackedBTreeEntry( + const std::string& file_name, int32_t bucket, + const std::optional& external_path = std::nullopt) const { + constexpr int64_t kRowCount = 3; + PAIMON_ASSIGN_OR_RAISE( + PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create( + /*data_level=*/2, {{fmt::format("{}.data", file_name), kRowCount}})); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr source_meta_bytes, + source_meta.Serialize(mem_pool_)); + auto index_file = std::make_shared( + "btree", file_name, /*file_size=*/7, kRowCount, /*dv_ranges=*/std::nullopt, + external_path, + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/kRowCount - 1, + /*index_field_id=*/5, /*extra_field_ids=*/std::nullopt, + /*index_meta=*/nullptr, source_meta_bytes)); + BinaryRow partition = + CreateManifestEntry("partition-source.data", bucket, FileKind::Add()).Partition(); + return IndexManifestEntry(FileKind::Add(), partition, bucket, index_file); + } + + Result IndexFilePath(const IndexManifestEntry& entry) const { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr index_path_factory, + path_factory_->CreateIndexFileFactory(entry.partition, entry.bucket)); + return index_path_factory->ToPath(entry.index_file); + } + private: + std::unique_ptr dir_; std::string test_data_path_; std::vector partition_keys_; std::shared_ptr schema_; @@ -161,6 +227,7 @@ class ExpireSnapshotsTest : public testing::Test { std::shared_ptr executor_; std::shared_ptr manifest_list_; std::shared_ptr manifest_file_; + std::shared_ptr index_manifest_file_; std::shared_ptr fs_; std::shared_ptr path_factory_; }; @@ -169,32 +236,36 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { auto mgr = std::make_shared(fs_, test_data_path_); { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_OK_AND_ASSIGN(int32_t count, expire.Expire()); ASSERT_EQ(count, 0); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "0"}})); - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "9"}})); - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_OK_AND_ASSIGN(int32_t count, expire.Expire()); ASSERT_EQ(count, 0); } @@ -203,26 +274,227 @@ TEST_F(ExpireSnapshotsTest, TestInvalidInput) { CoreOptions::FromMap({{Options::SNAPSHOT_EXPIRE_LIMIT, "-1"}, {Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "10"}, {Options::SNAPSHOT_NUM_RETAINED_MAX, "10"}})); - ExpireSnapshots expire(nullptr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(nullptr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); ASSERT_NOK(expire.Expire()); } } +TEST_F(ExpireSnapshotsTest, TestExpireSourceBackedBTreeIndexFiles) { + auto mgr = std::make_shared(fs_, test_data_path_); + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, index_manifest_file_, + fs_, options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + + ASSERT_OK_AND_ASSIGN(IndexManifestEntry shared_index, + CreateSourceBackedBTreeEntry("shared-btree.index", /*bucket=*/0)); + std::string expired_external_path = test_data_path_ + "/external-expired-btree.index"; + ASSERT_OK_AND_ASSIGN( + IndexManifestEntry expired_index, + CreateSourceBackedBTreeEntry("expired-btree.index", /*bucket=*/0, expired_external_path)); + ASSERT_OK_AND_ASSIGN(IndexManifestEntry retained_index, + CreateSourceBackedBTreeEntry("retained-btree.index", /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::string shared_path, IndexFilePath(shared_index)); + ASSERT_OK_AND_ASSIGN(std::string expired_path, IndexFilePath(expired_index)); + ASSERT_OK_AND_ASSIGN(std::string retained_path, IndexFilePath(retained_index)); + ASSERT_OK(fs_->WriteFile(shared_path, "payload", /*overwrite=*/false)); + ASSERT_OK(fs_->WriteFile(expired_path, "payload", /*overwrite=*/false)); + ASSERT_OK(fs_->WriteFile(retained_path, "payload", /*overwrite=*/false)); + + ASSERT_OK_AND_ASSIGN( + std::optional expired_manifest, + index_manifest_file_->WriteIndexFiles(std::nullopt, {shared_index, expired_index})); + ASSERT_TRUE(expired_manifest); + ASSERT_OK_AND_ASSIGN( + std::optional retained_manifest, + index_manifest_file_->WriteIndexFiles(std::nullopt, {shared_index, retained_index})); + ASSERT_TRUE(retained_manifest); + + std::set skipping_set; + ASSERT_OK(expire.AddIndexManifestToSkippingSet(retained_manifest, &skipping_set)); + ASSERT_GT(skipping_set.count(retained_manifest.value()), 0); + ASSERT_GT(skipping_set.count(shared_index.index_file->FileName()), 0); + ASSERT_GT(skipping_set.count(retained_index.index_file->FileName()), 0); + + ASSERT_OK(expire.CleanUnusedIndexManifest(expired_manifest, &skipping_set)); + ASSERT_OK_AND_ASSIGN(bool shared_exists, fs_->Exists(shared_path)); + ASSERT_TRUE(shared_exists); + ASSERT_OK_AND_ASSIGN(bool expired_exists, fs_->Exists(expired_path)); + ASSERT_FALSE(expired_exists); + ASSERT_OK_AND_ASSIGN(bool retained_exists, fs_->Exists(retained_path)); + ASSERT_TRUE(retained_exists); + ASSERT_OK_AND_ASSIGN(bool expired_manifest_exists, + fs_->Exists(path_factory_->ToManifestFilePath(expired_manifest.value()))); + ASSERT_FALSE(expired_manifest_exists); + ASSERT_OK_AND_ASSIGN(bool retained_manifest_exists, + fs_->Exists(path_factory_->ToManifestFilePath(retained_manifest.value()))); + ASSERT_TRUE(retained_manifest_exists); + + ASSERT_OK( + expire.CleanUnusedIndexManifest(std::string("missing-index-manifest"), &skipping_set)); + std::set missing_skipping_set; + ASSERT_NOK(expire.AddIndexManifestToSkippingSet(std::string("missing-index-manifest"), + &missing_skipping_set)); +} + +TEST_F(ExpireSnapshotsTest, TestExternalIndexDeleteFailureKeepsManifestForRetry) { + auto mgr = std::make_shared(fs_, test_data_path_); + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + + std::string external_path = test_data_path_ + "/retry-external-btree.index"; + ASSERT_OK_AND_ASSIGN( + IndexManifestEntry external_index, + CreateSourceBackedBTreeEntry("retry-btree.index", /*bucket=*/0, external_path)); + ASSERT_OK(fs_->WriteFile(external_path, "payload", /*overwrite=*/false)); + ASSERT_OK_AND_ASSIGN(std::optional index_manifest, + index_manifest_file_->WriteIndexFiles(std::nullopt, {external_index})); + ASSERT_TRUE(index_manifest); + + auto failing_fs = std::make_shared(external_path); + ExpireSnapshots failing_expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, failing_fs, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); + std::set skipping_set; + ASSERT_NOK_WITH_MSG(failing_expire.CleanUnusedIndexManifest(index_manifest, &skipping_set), + "injected delete failure"); + ASSERT_OK_AND_ASSIGN(bool payload_exists, fs_->Exists(external_path)); + ASSERT_TRUE(payload_exists); + std::string manifest_path = path_factory_->ToManifestFilePath(index_manifest.value()); + ASSERT_OK_AND_ASSIGN(bool manifest_exists, fs_->Exists(manifest_path)); + ASSERT_TRUE(manifest_exists); + + ExpireSnapshots retry_expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); + ASSERT_OK(retry_expire.CleanUnusedIndexManifest(index_manifest, &skipping_set)); + ASSERT_OK_AND_ASSIGN(payload_exists, fs_->Exists(external_path)); + ASSERT_FALSE(payload_exists); + ASSERT_OK_AND_ASSIGN(manifest_exists, fs_->Exists(manifest_path)); + ASSERT_FALSE(manifest_exists); +} + +TEST_F(ExpireSnapshotsTest, TestRejectExpirationWhileAnotherBranchExists) { + auto mgr = std::make_shared(fs_, test_data_path_); + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, index_manifest_file_, + fs_, options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ASSERT_OK(fs_->Mkdirs(PathUtil::JoinPath(test_data_path_, "branch/branch-dev"))); + ASSERT_NOK_WITH_MSG(expire.ExpireUntil(/*earliest_snapshot_id=*/1, + /*end_exclusive_id=*/2), + "cross-branch file retention is not supported"); +} + +TEST_F(ExpireSnapshotsTest, TestExpireKeepsSourceBackedBTreeIndexReferencedByTag) { + auto mgr = std::make_shared(fs_, test_data_path_); + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::SNAPSHOT_NUM_RETAINED_MIN, "1"}, + {Options::SNAPSHOT_NUM_RETAINED_MAX, "1"}})); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, index_manifest_file_, + fs_, options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + + ASSERT_OK_AND_ASSIGN(IndexManifestEntry tagged_index, + CreateSourceBackedBTreeEntry("tagged-btree.index", /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::string tagged_index_path, IndexFilePath(tagged_index)); + ASSERT_OK(fs_->WriteFile(tagged_index_path, "payload", /*overwrite=*/false)); + ASSERT_OK_AND_ASSIGN(std::optional tagged_index_manifest, + index_manifest_file_->WriteIndexFiles(std::nullopt, {tagged_index})); + ASSERT_TRUE(tagged_index_manifest); + + using ManifestListMeta = std::pair; + ManifestEntry tagged_data = CreateManifestEntry("tagged.data", /*bucket=*/0, FileKind::Add()); + ASSERT_OK_AND_ASSIGN(std::string tagged_bucket_path, + path_factory_->BucketPath(tagged_data.Partition(), tagged_data.Bucket())); + ASSERT_OK(fs_->Mkdirs(tagged_bucket_path)); + std::string tagged_data_path = PathUtil::JoinPath(tagged_bucket_path, tagged_data.FileName()); + ASSERT_OK(fs_->WriteFile(tagged_data_path, "data", /*overwrite=*/false)); + ASSERT_OK_AND_ASSIGN(std::vector tagged_data_manifests, + manifest_file_->Write({tagged_data})); + ASSERT_OK_AND_ASSIGN(ManifestListMeta tagged_base_manifest_list, + manifest_list_->Write(tagged_data_manifests)); + ASSERT_OK_AND_ASSIGN(ManifestListMeta tagged_delta_manifest_list, manifest_list_->Write({})); + ASSERT_OK_AND_ASSIGN(ManifestListMeta retained_base_manifest_list, manifest_list_->Write({})); + ManifestEntry deleted_tagged_data = + CreateManifestEntry(tagged_data.FileName(), /*bucket=*/0, FileKind::Delete()); + ASSERT_OK_AND_ASSIGN(std::vector retained_data_manifests, + manifest_file_->Write({deleted_tagged_data})); + ASSERT_OK_AND_ASSIGN(ManifestListMeta retained_delta_manifest_list, + manifest_list_->Write(retained_data_manifests)); + + Snapshot tagged_snapshot( + /*id=*/1, /*schema_id=*/0, tagged_base_manifest_list.first, + tagged_base_manifest_list.second, tagged_delta_manifest_list.first, + tagged_delta_manifest_list.second, /*changelog_manifest_list=*/std::nullopt, + /*changelog_manifest_list_size=*/std::nullopt, tagged_index_manifest, + /*commit_user=*/"test", /*commit_identifier=*/1, Snapshot::CommitKind::Append(), + /*time_millis=*/0, /*total_record_count=*/3, /*delta_record_count=*/3, + /*changelog_record_count=*/std::nullopt, /*watermark=*/std::nullopt, + /*statistics=*/std::nullopt, /*properties=*/std::nullopt, /*next_row_id=*/std::nullopt); + Snapshot retained_snapshot( + /*id=*/2, /*schema_id=*/0, retained_base_manifest_list.first, + retained_base_manifest_list.second, retained_delta_manifest_list.first, + retained_delta_manifest_list.second, /*changelog_manifest_list=*/std::nullopt, + /*changelog_manifest_list_size=*/std::nullopt, /*index_manifest=*/std::nullopt, + /*commit_user=*/"test", /*commit_identifier=*/2, Snapshot::CommitKind::Append(), + /*time_millis=*/0, /*total_record_count=*/3, /*delta_record_count=*/0, + /*changelog_record_count=*/std::nullopt, /*watermark=*/std::nullopt, + /*statistics=*/std::nullopt, /*properties=*/std::nullopt, /*next_row_id=*/std::nullopt); + + ASSERT_OK(fs_->Mkdirs(mgr->SnapshotDirectory())); + ASSERT_OK_AND_ASSIGN(std::string tagged_snapshot_json, tagged_snapshot.ToJsonString()); + ASSERT_OK(fs_->WriteFile(mgr->SnapshotPath(tagged_snapshot.Id()), tagged_snapshot_json, + /*overwrite=*/false)); + ASSERT_OK_AND_ASSIGN(std::string retained_snapshot_json, retained_snapshot.ToJsonString()); + ASSERT_OK(fs_->WriteFile(mgr->SnapshotPath(retained_snapshot.Id()), retained_snapshot_json, + /*overwrite=*/false)); + + Tag tag(tagged_snapshot.Version(), tagged_snapshot.Id(), tagged_snapshot.SchemaId(), + tagged_snapshot.BaseManifestList(), tagged_snapshot.BaseManifestListSize(), + tagged_snapshot.DeltaManifestList(), tagged_snapshot.DeltaManifestListSize(), + tagged_snapshot.ChangelogManifestList(), tagged_snapshot.ChangelogManifestListSize(), + tagged_snapshot.IndexManifest(), tagged_snapshot.CommitUser(), + tagged_snapshot.CommitIdentifier(), tagged_snapshot.GetCommitKind(), + tagged_snapshot.TimeMillis(), tagged_snapshot.TotalRecordCount(), + tagged_snapshot.DeltaRecordCount(), tagged_snapshot.ChangelogRecordCount(), + tagged_snapshot.Watermark(), tagged_snapshot.Statistics(), tagged_snapshot.Properties(), + tagged_snapshot.NextRowId(), /*tag_create_time=*/std::nullopt, + /*tag_time_retained=*/std::nullopt); + TagManager tag_manager(fs_, mgr->RootPath(), mgr->Branch()); + ASSERT_OK(fs_->Mkdirs(tag_manager.TagDirectory())); + ASSERT_OK_AND_ASSIGN(std::string tag_json, tag.ToJsonString()); + ASSERT_OK(fs_->WriteFile(tag_manager.TagPath("tagged"), tag_json, /*overwrite=*/false)); + + ASSERT_OK_AND_ASSIGN(int32_t expired_count, expire.Expire()); + ASSERT_EQ(expired_count, 1); + ASSERT_OK_AND_ASSIGN(bool expired_snapshot_exists, + fs_->Exists(mgr->SnapshotPath(tagged_snapshot.Id()))); + ASSERT_FALSE(expired_snapshot_exists); + ASSERT_OK_AND_ASSIGN( + bool tagged_index_manifest_exists, + fs_->Exists(path_factory_->ToManifestFilePath(tagged_index_manifest.value()))); + ASSERT_TRUE(tagged_index_manifest_exists); + ASSERT_OK_AND_ASSIGN(bool tagged_index_exists, fs_->Exists(tagged_index_path)); + ASSERT_TRUE(tagged_index_exists); + ASSERT_OK_AND_ASSIGN(bool tagged_data_exists, fs_->Exists(tagged_data_path)); + ASSERT_TRUE(tagged_data_exists); +} + TEST_F(ExpireSnapshotsTest, TestGetDataFileToDelete) { auto mgr = std::make_shared(fs_, test_data_path_); ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({})); { - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); std::map data_file_to_delete; std::vector data_file_entries; data_file_entries.push_back(CreateManifestEntry("file1", /*bucket=*/0, FileKind::Delete())); @@ -236,8 +508,9 @@ TEST_F(ExpireSnapshotsTest, TestGetDataFileToDelete) { {test_data_path_ + "/f0=true/f3=3/bucket-2/file3", data_file_entries[3]}})); } { - ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, fs_, - options.GetExpireConfig(), options.RealtimeEnabled(), executor_); + ExpireSnapshots expire(mgr, path_factory_, manifest_list_, manifest_file_, + index_manifest_file_, fs_, options.GetExpireConfig(), + options.RealtimeEnabled(), executor_); std::map data_file_to_delete; std::vector data_file_entries; data_file_entries.push_back(CreateManifestEntry("file1", /*bucket=*/0, FileKind::Add())); diff --git a/src/paimon/core/operation/file_store_commit.cpp b/src/paimon/core/operation/file_store_commit.cpp index ad0942ade..df836d134 100644 --- a/src/paimon/core/operation/file_store_commit.cpp +++ b/src/paimon/core/operation/file_store_commit.cpp @@ -174,8 +174,9 @@ Result> FileStoreCommit::Create( options.GetBucket(), ctx->GetMemoryPool(), options)); auto expire_snapshots = std::make_shared( - snapshot_manager, path_factory, manifest_list, manifest_file, options.GetFileSystem(), - options.GetExpireConfig(), options.RealtimeEnabled(), ctx->GetExecutor()); + snapshot_manager, path_factory, manifest_list, manifest_file, index_manifest_file, + options.GetFileSystem(), options.GetExpireConfig(), options.RealtimeEnabled(), + ctx->GetExecutor()); CommitScanner::ScanSupplier scan_supplier; if (table_schema.value()->PrimaryKeys().empty()) { diff --git a/src/paimon/core/operation/file_store_commit_impl.cpp b/src/paimon/core/operation/file_store_commit_impl.cpp index 8f2b3c732..5c1552e8f 100644 --- a/src/paimon/core/operation/file_store_commit_impl.cpp +++ b/src/paimon/core/operation/file_store_commit_impl.cpp @@ -95,7 +95,6 @@ namespace { constexpr const char* kCommitStrictModeLastSafeSnapshot = "commit.strict-mode.last-safe-snapshot"; constexpr const char* kSequenceSnapshotOrdering = "sequence.snapshot-ordering"; -constexpr const char* kPkClusteringOverride = "pk-clustering-override"; } // namespace @@ -109,8 +108,8 @@ Status FileStoreCommitImpl::ValidateCommitOptions(const CoreOptions& options) { if (raw_options.find(kSequenceSnapshotOrdering) != raw_options.end()) { unsupported_options.emplace_back(kSequenceSnapshotOrdering); } - if (raw_options.find(kPkClusteringOverride) != raw_options.end()) { - unsupported_options.emplace_back(kPkClusteringOverride); + if (raw_options.find(Options::PK_CLUSTERING_OVERRIDE) != raw_options.end()) { + unsupported_options.emplace_back(Options::PK_CLUSTERING_OVERRIDE); } if (!unsupported_options.empty()) { diff --git a/src/paimon/core/operation/file_store_commit_impl_test.cpp b/src/paimon/core/operation/file_store_commit_impl_test.cpp index b1b8677f8..6490e3726 100644 --- a/src/paimon/core/operation/file_store_commit_impl_test.cpp +++ b/src/paimon/core/operation/file_store_commit_impl_test.cpp @@ -2538,7 +2538,7 @@ TEST_F(FileStoreCommitImplTest, TestFixedBucketPKTableCommitAllowed) { TEST_F(FileStoreCommitImplTest, ValidateCommitOptionsRejectsUnsupportedOptions) { const std::vector unsupported_keys = {"commit.strict-mode.last-safe-snapshot", "sequence.snapshot-ordering", - "pk-clustering-override"}; + Options::PK_CLUSTERING_OVERRIDE}; for (const auto& key : unsupported_keys) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{key, "true"}})); ASSERT_NOK_WITH_MSG(FileStoreCommitImpl::ValidateCommitOptions(options), diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index fa1294d1b..ed0457b32 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -18,15 +18,22 @@ #include "paimon/file_store_write.h" +#include #include #include #include +#include #include "fmt/format.h" +#include "paimon/common/global_index/btree/btree_defs.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/string_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/disk/io_manager.h" +#include "paimon/core/index/index_file_handler.h" +#include "paimon/core/index/pk/bucketed_primary_key_index_maintainer.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" #include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/mergetree/compact/lookup_merge_function.h" #include "paimon/core/mergetree/compact/merge_function.h" @@ -59,6 +66,28 @@ class MergeFunctionWrapper; namespace { +Result HasHistoricalPrimaryKeyBTreeDefinition( + const std::shared_ptr& schema_manager, int64_t current_schema_id) { + if (current_schema_id <= TableSchema::FIRST_SCHEMA_ID) { + return false; + } + PAIMON_ASSIGN_OR_RAISE(std::vector schema_ids, schema_manager->ListAllIds()); + for (int64_t schema_id : schema_ids) { + if (schema_id >= current_schema_id) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr historical_schema, + schema_manager->ReadSchema(schema_id)); + const auto& historical_options = historical_schema->Options(); + auto option = historical_options.find(Options::PK_BTREE_INDEX_COLUMNS); + if (option != historical_options.end() && + !StringUtils::IsNullOrWhitespaceOnly(option->second)) { + return true; + } + } + return false; +} + Status RestoreRealtimeCommittedProgress(const std::shared_ptr& realtime_context, const std::shared_ptr& snapshot_manager, const CoreOptions& options) { @@ -255,27 +284,73 @@ Result> FileStoreWrite::Create(std::unique_ptrFields(), options)); std::shared_ptr dv_maintainer_factory; - if (options.DeletionVectorsEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions primary_key_index_definitions, + PrimaryKeyIndexDefinitions::Create(*schema)); + bool has_btree_index = false; + for (const PrimaryKeyIndexDefinition& definition : + primary_key_index_definitions.Definitions()) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE) { + has_btree_index = true; + break; + } + } + std::optional latest_snapshot; + bool has_existing_index_manifest = false; + if (!has_btree_index) { + PAIMON_ASSIGN_OR_RAISE( + bool has_historical_btree_index, + HasHistoricalPrimaryKeyBTreeDefinition(schema_manager, schema->Id())); + if (has_historical_btree_index) { + PAIMON_ASSIGN_OR_RAISE(latest_snapshot, snapshot_manager->LatestSnapshot()); + has_existing_index_manifest = + latest_snapshot.has_value() && latest_snapshot->IndexManifest().has_value(); + } + } + std::shared_ptr index_file_handler; + if (options.DeletionVectorsEnabled() || has_btree_index || has_existing_index_manifest) { PAIMON_ASSIGN_OR_RAISE( std::unique_ptr index_manifest_file, IndexManifestFile::Create(options.GetFileSystem(), options.GetManifestFormat(), options.GetManifestCompression(), file_store_path_factory, options.GetBucket(), ctx->GetMemoryPool(), options)); - auto index_file_handler = std::make_shared( + index_file_handler = std::make_shared( options.GetFileSystem(), std::move(index_manifest_file), std::make_shared(file_store_path_factory), options.DeletionVectorsBitmap64(), ctx->GetMemoryPool()); + } + bool has_restored_btree_index = false; + if (!has_btree_index && has_existing_index_manifest) { + PAIMON_ASSIGN_OR_RAISE( + std::vector restored_btree_entries, + index_file_handler->Scan( + latest_snapshot.value(), [](const IndexManifestEntry& entry) -> bool { + return entry.index_file->IndexType() == BtreeDefs::kIdentifier && + IndexFileHandler::IsPrimaryKeySourceIndex(*entry.index_file); + })); + has_restored_btree_index = !restored_btree_entries.empty(); + } + if (options.DeletionVectorsEnabled()) { dv_maintainer_factory = std::make_shared(index_file_handler); } + std::shared_ptr + primary_key_index_maintainer_factory; + if (has_btree_index || has_restored_btree_index) { + PAIMON_ASSIGN_OR_RAISE( + primary_key_index_maintainer_factory, + BucketedPrimaryKeyIndexMaintainer::Factory::Create( + ctx->GetRootPath(), branch, schema, primary_key_index_definitions.Definitions(), + file_store_path_factory, index_file_handler, options, io_manager, + ctx->EnableMultiThreadSpill(), ctx->GetExecutor(), ctx->GetMemoryPool())); + } return std::make_unique( file_store_path_factory, snapshot_manager, schema_manager, ctx->GetCommitUser(), ctx->GetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory, - io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options, - ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), - ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), - ctx->GetMemoryPool()); + primary_key_index_maintainer_factory, io_manager, key_comparator, + sequence_fields_comparator, merge_function_wrapper, options, ignore_previous_files, + ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), ctx->EnableMultiThreadSpill(), + ctx->GetRealtimeContext(), ctx->GetExecutor(), ctx->GetMemoryPool()); } } diff --git a/src/paimon/core/operation/file_system_write_restore.h b/src/paimon/core/operation/file_system_write_restore.h index 630ea83dd..7a23249b0 100644 --- a/src/paimon/core/operation/file_system_write_restore.h +++ b/src/paimon/core/operation/file_system_write_restore.h @@ -39,7 +39,7 @@ class FileSystemWriteRestore : public WriteRestore { /// /// @param snapshot_manager Snapshot manager used to locate restore state. /// @param scan Scan used to load existing files. - /// @param index_file_handler Handler used to restore deletion-vector indexes. + /// @param index_file_handler Handler used to restore index files. FileSystemWriteRestore(const std::shared_ptr& snapshot_manager, std::unique_ptr&& scan, const std::shared_ptr& index_file_handler) @@ -58,8 +58,8 @@ class FileSystemWriteRestore : public WriteRestore { } Result> GetRestoreFiles( - const BinaryRow& partition, int32_t bucket, - bool scan_deletion_vectors_index) const override { + const BinaryRow& partition, int32_t bucket, bool scan_deletion_vectors_index, + bool scan_source_index_payloads) const override { // TODO(yonghao.fyh): java paimon doesn't use snapshot_manager.LatestSnapshot() here, // because they don't want to flood the catalog with high concurrency PAIMON_ASSIGN_OR_RAISE(std::optional snapshot, @@ -84,9 +84,19 @@ class FileSystemWriteRestore : public WriteRestore { partition, bucket)); } + std::vector> primary_key_index_payloads; + if (scan_source_index_payloads) { + if (index_file_handler_ == nullptr) { + return Status::Invalid("Primary-key index restore requires an index file handler."); + } + PAIMON_ASSIGN_OR_RAISE( + primary_key_index_payloads, + index_file_handler_->ScanPrimaryKeyIndexes(snapshot.value(), partition, bucket)); + } + return std::make_shared(snapshot, total_buckets, restore_data_files, /*dynamic_bucket_index=*/nullptr, - deletion_vectors_index); + deletion_vectors_index, primary_key_index_payloads); } private: diff --git a/src/paimon/core/operation/file_system_write_restore_test.cpp b/src/paimon/core/operation/file_system_write_restore_test.cpp index 4740e64db..d4e8a9196 100644 --- a/src/paimon/core/operation/file_system_write_restore_test.cpp +++ b/src/paimon/core/operation/file_system_write_restore_test.cpp @@ -20,13 +20,86 @@ #include #include +#include +#include +#include #include "gtest/gtest.h" +#include "paimon/common/types/data_field.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/index_file_handler.h" +#include "paimon/core/manifest/index_manifest_file.h" +#include "paimon/core/manifest/manifest_file.h" +#include "paimon/core/manifest/manifest_list.h" +#include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/utils/field_mapping.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/core/utils/index_file_path_factories.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/executor.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/scan_context.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +Result> CreateRestore( + const std::string& table_path, const std::shared_ptr& pool) { + auto fs = std::make_shared(); + auto schema_manager = std::make_shared(fs, table_path); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr table_schema, + schema_manager->ReadSchema(/*schema_id=*/0)); + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(table_schema->Options())); + std::shared_ptr arrow_schema = + DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); + PAIMON_ASSIGN_OR_RAISE(std::vector external_paths, options.CreateExternalPaths()); + PAIMON_ASSIGN_OR_RAISE(std::optional global_index_external_path, + options.CreateGlobalIndexExternalPath()); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr path_factory, + FileStorePathFactory::Create( + table_path, arrow_schema, table_schema->PartitionKeys(), + options.GetPartitionDefaultName(), options.GetFileFormat()->Identifier(), + options.DataFilePrefix(), options.LegacyPartitionNameEnabled(), external_paths, + global_index_external_path, options.IndexFileInDataFileDir(), pool)); + auto snapshot_manager = std::make_shared(fs, table_path); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr manifest_list, + ManifestList::Create(fs, options.GetManifestFormat(), options.GetManifestCompression(), + path_factory, options.GetCache(), pool)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr partition_schema, + FieldMapping::GetPartitionSchema(arrow_schema, table_schema->PartitionKeys())); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr manifest_file, + ManifestFile::Create(fs, options.GetManifestFormat(), options.GetManifestCompression(), + path_factory, options.GetManifestTargetFileSize(), pool, options, + partition_schema)); + auto scan_filter = std::make_shared( + /*predicate=*/nullptr, + /*partition_filters=*/std::vector>{}, + /*bucket_filter=*/std::nullopt); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr scan, + KeyValueFileStoreScan::Create(snapshot_manager, schema_manager, manifest_list, + manifest_file, table_schema, arrow_schema, scan_filter, + options, CreateDefaultExecutor(), pool)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr index_manifest_file, + IndexManifestFile::Create(fs, options.GetManifestFormat(), options.GetManifestCompression(), + path_factory, options.GetBucket(), pool, options)); + auto index_file_handler = std::make_shared( + fs, std::move(index_manifest_file), std::make_shared(path_factory), + options.DeletionVectorsBitmap64(), pool); + return std::make_unique(snapshot_manager, std::move(scan), + index_file_handler); +} + +} // namespace TEST(FileSystemWriteRestoreTest, LatestCommittedIdentifierNoSnapshot) { auto fs = std::make_shared(); @@ -64,11 +137,29 @@ TEST(FileSystemWriteRestoreTest, GetRestoreFilesReturnsEmptyWhenNoLatestSnapshot ASSERT_OK_AND_ASSIGN(std::shared_ptr files, restore.GetRestoreFiles(BinaryRow::EmptyRow(), /*bucket=*/0, - /*scan_deletion_vectors_index=*/true)); + /*scan_deletion_vectors_index=*/true, + /*scan_source_index_payloads=*/false)); ASSERT_FALSE(files->GetSnapshot().has_value()); ASSERT_FALSE(files->TotalBuckets().has_value()); ASSERT_TRUE(files->DataFiles().empty()); ASSERT_TRUE(files->DeleteVectorsIndex().empty()); } +TEST(FileSystemWriteRestoreTest, RestoresSourceBackedIndexPayloads) { + const std::string table_path = + paimon::test::GetDataDir() + "/orc/pk_btree_source_meta.db/pk_btree_source_meta/"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr restore, + CreateRestore(table_path, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr files, + restore->GetRestoreFiles(BinaryRow::EmptyRow(), /*bucket=*/0, + /*scan_deletion_vectors_index=*/false, + /*scan_source_index_payloads=*/true)); + + ASSERT_EQ(files->PrimaryKeyIndexPayloads().size(), 1); + const std::optional& global_index_meta = + files->PrimaryKeyIndexPayloads()[0]->GetGlobalIndexMeta(); + ASSERT_TRUE(global_index_meta.has_value()); + ASSERT_NE(global_index_meta->source_meta, nullptr); +} + } // namespace paimon::test diff --git a/src/paimon/core/operation/key_value_file_store_write.cpp b/src/paimon/core/operation/key_value_file_store_write.cpp index e8e8cd3d9..20765bef8 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -63,6 +63,8 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr& schema, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, + const std::shared_ptr& + primary_key_index_maintainer_factory, const std::shared_ptr& io_manager, const std::shared_ptr& key_comparator, const std::shared_ptr& user_defined_seq_comparator, @@ -71,11 +73,11 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( bool ignore_num_bucket_check, bool enable_multi_thread_spill, const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool) - : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, - root_path, table_schema, schema, /*write_schema=*/schema, - partition_schema, dv_maintainer_factory, io_manager, options, - ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, - executor, pool), + : AbstractFileStoreWrite( + file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, + table_schema, schema, /*write_schema=*/schema, partition_schema, dv_maintainer_factory, + primary_key_index_maintainer_factory, io_manager, options, ignore_previous_files, + is_streaming_mode, ignore_num_bucket_check, executor, pool), enable_multi_thread_spill_(enable_multi_thread_spill), realtime_context_(realtime_context), key_comparator_(key_comparator), diff --git a/src/paimon/core/operation/key_value_file_store_write.h b/src/paimon/core/operation/key_value_file_store_write.h index 66c362f2e..45afe9225 100644 --- a/src/paimon/core/operation/key_value_file_store_write.h +++ b/src/paimon/core/operation/key_value_file_store_write.h @@ -60,6 +60,8 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr& schema, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, + const std::shared_ptr& + primary_key_index_maintainer_factory, const std::shared_ptr& io_manager, const std::shared_ptr& key_comparator, const std::shared_ptr& user_defined_seq_comparator, diff --git a/src/paimon/core/operation/orphan_files_cleaner.cpp b/src/paimon/core/operation/orphan_files_cleaner.cpp index ac69d5a37..397ba68ea 100644 --- a/src/paimon/core/operation/orphan_files_cleaner.cpp +++ b/src/paimon/core/operation/orphan_files_cleaner.cpp @@ -27,6 +27,7 @@ #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/core/core_options.h" +#include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/operation/orphan_files_cleaner_impl.h" @@ -169,9 +170,6 @@ Result> OrphanFilesCleaner::Create( if (table_schema == std::nullopt) { return Status::Invalid("not found latest schema"); } - if (!table_schema.value()->PrimaryKeys().empty()) { - return Status::NotImplemented("orphan files cleaner only support append table"); - } // merge options const auto& schema = table_schema.value(); auto opts = schema->Options(); @@ -208,10 +206,15 @@ Result> OrphanFilesCleaner::Create( options.GetManifestCompression(), path_factory, options.GetManifestTargetFileSize(), ctx->GetMemoryPool(), options, partition_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr index_manifest_file, + IndexManifestFile::Create(options.GetFileSystem(), options.GetManifestFormat(), + options.GetManifestCompression(), path_factory, + options.GetBucket(), ctx->GetMemoryPool(), options)); return std::make_unique( ctx->GetMemoryPool(), ctx->GetExecutor(), arrow_schema, ctx->GetRootPath(), options, snapshot_manager, schema->PartitionKeys(), manifest_file, manifest_list, - ctx->GetOlderThanMs(), ctx->GetFileRetainCondition()); + index_manifest_file, ctx->GetOlderThanMs(), ctx->GetFileRetainCondition()); } } // namespace paimon diff --git a/src/paimon/core/operation/orphan_files_cleaner_impl.cpp b/src/paimon/core/operation/orphan_files_cleaner_impl.cpp index 015d13408..64d38c13a 100644 --- a/src/paimon/core/operation/orphan_files_cleaner_impl.cpp +++ b/src/paimon/core/operation/orphan_files_cleaner_impl.cpp @@ -31,6 +31,9 @@ #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/string_utils.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/manifest/index_manifest_file.h" #include "paimon/core/manifest/manifest_entry.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_file_meta.h" @@ -54,7 +57,8 @@ OrphanFilesCleanerImpl::OrphanFilesCleanerImpl( const CoreOptions& options, const std::shared_ptr& snapshot_manager, const std::vector& partition_keys, const std::shared_ptr& manifest_file, - const std::shared_ptr& manifest_list, int64_t older_than_ms, + const std::shared_ptr& manifest_list, + const std::shared_ptr& index_manifest_file, int64_t older_than_ms, std::function should_be_retained) : memory_pool_(memory_pool), executor_(executor), @@ -66,13 +70,18 @@ OrphanFilesCleanerImpl::OrphanFilesCleanerImpl( partition_keys_(partition_keys), manifest_file_(manifest_file), manifest_list_(manifest_list), + index_manifest_file_(index_manifest_file), older_than_ms_(older_than_ms), should_be_retained_(should_be_retained), metrics_(std::make_shared()) {} bool OrphanFilesCleanerImpl::SupportToClean(const std::string& file_name) { static std::vector> supported_pattern = { - {"manifest-", ""}, {"manifest-list-", ""}, {".", ".tmp"}}; + {"manifest-", ""}, + {"manifest-list-", ""}, + {"index-manifest-", ""}, + {"index-", ""}, + {".", ".tmp"}}; for (const auto& pattern : supported_pattern) { if (StringUtils::StartsWith(file_name, pattern.first) && StringUtils::EndsWith(file_name, pattern.second)) { @@ -86,7 +95,9 @@ bool OrphanFilesCleanerImpl::SupportToClean(const std::string& file_name) { return true; } } - return StringUtils::EndsWith(file_name, ".offsets"); + return StringUtils::EndsWith(file_name, ".offsets") || + (file_name.find("-global-index-") != std::string::npos && + StringUtils::EndsWith(file_name, ".index")); } Result> OrphanFilesCleanerImpl::Clean() { @@ -157,12 +168,12 @@ Result> OrphanFilesCleanerImpl::ListPaimonFileDirs() const std::set paimon_file_dirs; paimon_file_dirs.insert(snapshot_manager_->SnapshotDirectory()); paimon_file_dirs.insert(FileStorePathFactory::ManifestPath(root_path_)); + paimon_file_dirs.insert(FileStorePathFactory::IndexPath(root_path_)); if (options_.RealtimeEnabled()) { paimon_file_dirs.insert( RealtimeCommitProperties::OffsetsDirectory(root_path_, options_.GetBranch())); } - // TODO(jinli.zjw): support clean index, stats, changelog in the future - // paimon_file_dirs.insert(FileStorePathFactory::IndexPath(root_path_)); + // TODO(jinli.zjw): support clean stats and changelog in the future // paimon_file_dirs.insert(FileStorePathFactory::StatisticsPath(root_path_)); std::set file_dirs = ListFileDirs(root_path_, partition_keys_.size()); paimon_file_dirs.insert(file_dirs.begin(), file_dirs.end()); @@ -173,14 +184,6 @@ Result> OrphanFilesCleanerImpl::ListPaimonFileDirs() const return Status::Invalid( "OrphanFilesCleaner do not support cleaning table with external paths"); } - // TODO(liancheng): support clean external paths in the future - // PAIMON_ASSIGN_OR_RAISE(std::vector external_paths, - // options_.CreateExternalPaths()); - // for (const auto& external_path : external_paths) { - // std::set external_file_dirs = - // ListFileDirs(external_path, partition_keys_.size()); - // paimon_file_dirs.insert(external_file_dirs.begin(), external_file_dirs.end()); - // } metrics_->SetCounter(CleanMetrics::CLEAN_LIST_DIRECTORIES_DURATION, duration.Get()); metrics_->SetCounter(CleanMetrics::CLEAN_LIST_DIRECTORIES, static_cast(paimon_file_dirs.size())); @@ -267,14 +270,6 @@ Result> OrphanFilesCleanerImpl::GetUsedFiles() const { used_files.insert(changelog_manifest_list.value()); return Status::NotImplemented("OrphanFilesCleaner do not support clean changelog"); } - const std::optional& index_manifest_name = snapshot.IndexManifest(); - if (index_manifest_name) { - return Status::NotImplemented( - "OrphanFilesCleaner do not support clean index manifest"); - // TODO(jinli.zjw): support IndexManifestEntry and add tests - // used_files.insert(index_manifest_name.value()); - } - used_files_futures.emplace_back(Via( executor_.get(), [this, snapshot] { return GetUsedFilesBySnapshot(snapshot); })); } @@ -299,6 +294,16 @@ Result> OrphanFilesCleanerImpl::GetUsedFilesBySnapshot( used_files.insert(SnapshotManager::SNAPSHOT_PREFIX + std::to_string(snapshot.Id())); used_files.insert(snapshot.BaseManifestList()); used_files.insert(snapshot.DeltaManifestList()); + const std::optional& index_manifest_name = snapshot.IndexManifest(); + if (index_manifest_name) { + used_files.insert(index_manifest_name.value()); + std::vector index_entries; + PAIMON_RETURN_NOT_OK(index_manifest_file_->ReadIfFileExist( + index_manifest_name.value(), /*filter=*/nullptr, &index_entries)); + for (const IndexManifestEntry& index_entry : index_entries) { + used_files.insert(index_entry.index_file->FileName()); + } + } if (options_.RealtimeEnabled()) { std::optional offsets_path = RealtimeCommitProperties::GetOffsetsPath(snapshot); @@ -319,6 +324,12 @@ Result> OrphanFilesCleanerImpl::GetUsedFilesBySnapshot( manifest.FileName(), /*filter=*/nullptr, &manifest_entries)); for (const auto& manifest_entry : manifest_entries) { used_files.insert(manifest_entry.FileName()); + for (const std::optional& extra_file : + manifest_entry.File()->extra_files) { + if (extra_file) { + used_files.insert(extra_file.value()); + } + } } } diff --git a/src/paimon/core/operation/orphan_files_cleaner_impl.h b/src/paimon/core/operation/orphan_files_cleaner_impl.h index 88cde44e4..0c1164256 100644 --- a/src/paimon/core/operation/orphan_files_cleaner_impl.h +++ b/src/paimon/core/operation/orphan_files_cleaner_impl.h @@ -45,6 +45,7 @@ namespace paimon { class SnapshotManager; class FileStorePathFactory; class FileSystem; +class IndexManifestFile; class ManifestFile; class ManifestList; class Executor; @@ -60,6 +61,7 @@ class OrphanFilesCleanerImpl : public OrphanFilesCleaner { const std::vector& partition_keys, const std::shared_ptr& manifest_file, const std::shared_ptr& manifest_list, + const std::shared_ptr& index_manifest_file, int64_t older_than_ms, std::function should_be_retained); @@ -91,6 +93,7 @@ class OrphanFilesCleanerImpl : public OrphanFilesCleaner { std::vector partition_keys_; std::shared_ptr manifest_file_; std::shared_ptr manifest_list_; + std::shared_ptr index_manifest_file_; int64_t older_than_ms_; std::function should_be_retained_; diff --git a/src/paimon/core/operation/orphan_files_cleaner_test.cpp b/src/paimon/core/operation/orphan_files_cleaner_test.cpp index 7749b692b..b9ffb261a 100644 --- a/src/paimon/core/operation/orphan_files_cleaner_test.cpp +++ b/src/paimon/core/operation/orphan_files_cleaner_test.cpp @@ -18,19 +18,88 @@ #include "paimon/orphan_files_cleaner.h" +#include #include #include +#include +#include +#include +#include +#include #include +#include #include "gtest/gtest.h" +#include "paimon/common/data/binary_row.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/global_index_meta.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/index/index_path_factory.h" +#include "paimon/core/index/pk/primary_key_index_source_meta.h" +#include "paimon/core/manifest/file_kind.h" +#include "paimon/core/manifest/index_manifest_entry.h" +#include "paimon/core/manifest/index_manifest_file.h" +#include "paimon/core/manifest/manifest_entry.h" +#include "paimon/core/manifest/manifest_file.h" +#include "paimon/core/manifest/manifest_file_meta.h" +#include "paimon/core/manifest/manifest_list.h" #include "paimon/core/operation/orphan_files_cleaner_impl.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/snapshot.h" +#include "paimon/core/utils/file_store_path_factory.h" #include "paimon/defs.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +Result CreateSourceBackedBTreeEntry( + const std::string& file_name, const BinaryRow& partition, + const std::optional& external_path, const std::shared_ptr& pool) { + constexpr int64_t kRowCount = 3; + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexSourceMeta source_meta, + PrimaryKeyIndexSourceMeta::Create( + /*data_level=*/2, {{file_name + ".data", kRowCount}})); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr source_meta_bytes, source_meta.Serialize(pool)); + auto index_file = std::make_shared( + "btree", file_name, /*file_size=*/7, kRowCount, /*dv_ranges=*/std::nullopt, external_path, + GlobalIndexMeta(/*row_range_start=*/0, /*row_range_end=*/kRowCount - 1, + /*index_field_id=*/0, /*extra_field_ids=*/std::nullopt, + /*index_meta=*/nullptr, source_meta_bytes)); + return IndexManifestEntry(FileKind::Add(), partition, /*bucket=*/0, index_file); +} + +Snapshot WithIndexManifest(const Snapshot& snapshot, const std::string& index_manifest) { + return Snapshot(snapshot.Id(), snapshot.SchemaId(), snapshot.BaseManifestList(), + snapshot.BaseManifestListSize(), snapshot.DeltaManifestList(), + snapshot.DeltaManifestListSize(), snapshot.ChangelogManifestList(), + snapshot.ChangelogManifestListSize(), index_manifest, snapshot.CommitUser(), + snapshot.CommitIdentifier(), snapshot.GetCommitKind(), snapshot.TimeMillis(), + snapshot.TotalRecordCount(), snapshot.DeltaRecordCount(), + snapshot.ChangelogRecordCount(), snapshot.Watermark(), snapshot.Statistics(), + snapshot.Properties(), snapshot.NextRowId()); +} + +Snapshot WithDataManifests(const Snapshot& snapshot, const std::string& base_manifest_list, + const std::optional& base_manifest_list_size, + const std::string& delta_manifest_list, + const std::optional& delta_manifest_list_size) { + return Snapshot(snapshot.Id(), snapshot.SchemaId(), base_manifest_list, base_manifest_list_size, + delta_manifest_list, delta_manifest_list_size, snapshot.ChangelogManifestList(), + snapshot.ChangelogManifestListSize(), snapshot.IndexManifest(), + snapshot.CommitUser(), snapshot.CommitIdentifier(), snapshot.GetCommitKind(), + snapshot.TimeMillis(), snapshot.TotalRecordCount(), snapshot.DeltaRecordCount(), + snapshot.ChangelogRecordCount(), snapshot.Watermark(), snapshot.Statistics(), + snapshot.Properties(), snapshot.NextRowId()); +} + +} // namespace TEST(OrphanFilesCleanerTest, TestSupportToClean) { ASSERT_TRUE( @@ -55,8 +124,12 @@ TEST(OrphanFilesCleanerTest, TestSupportToClean) { "changelog-ce64d06d-c4cd-456b-a1b3-ae570042620f-0.parquet")); ASSERT_FALSE(OrphanFilesCleanerImpl::SupportToClean( "data-5515726b-0f0f-4556-a942-e795e9f94c4a-0.orc.index")); - ASSERT_FALSE( + ASSERT_TRUE( OrphanFilesCleanerImpl::SupportToClean("index-aa60193d-d7cd-434f-bc1a-c1adb210e1f7-0")); + ASSERT_TRUE(OrphanFilesCleanerImpl::SupportToClean( + "index-manifest-aa60193d-d7cd-434f-bc1a-c1adb210e1f7-0")); + ASSERT_TRUE(OrphanFilesCleanerImpl::SupportToClean( + "btree-global-index-aa60193d-d7cd-434f-bc1a-c1adb210e1f7.index")); ASSERT_FALSE( OrphanFilesCleanerImpl::SupportToClean("data-2d5ea1ea-77c1-47ff-bb87-19a509962a37-0.json")); ASSERT_FALSE(OrphanFilesCleanerImpl::SupportToClean( @@ -64,13 +137,83 @@ TEST(OrphanFilesCleanerTest, TestSupportToClean) { } TEST(OrphanFilesCleanerTest, TestPkTable) { - std::string table_path = + std::string test_data_path = paimon::test::GetDataDir() + "/orc/pk_table_with_mor.db/pk_table_with_mor/"; + auto dir = UniqueTestDirectory::Create(); + std::string table_path = dir->Str(); + ASSERT_TRUE(TestUtil::CopyDirectory(test_data_path, table_path)); + CleanContextBuilder clean_context_builder(table_path); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr clean_context, + clean_context_builder.AddOption(Options::FILE_SYSTEM, "local").WithOlderThanMs(0).Finish()); + ASSERT_OK_AND_ASSIGN(auto cleaner, OrphanFilesCleaner::Create(std::move(clean_context))); + ASSERT_OK_AND_ASSIGN(std::set cleaned_paths, cleaner->Clean()); + ASSERT_TRUE(cleaned_paths.empty()); +} + +TEST(OrphanFilesCleanerTest, TestRetainLiveDataFileExtraFiles) { + std::string test_data_path = + paimon::test::GetDataDir() + "/orc/pk_table_with_mor.db/pk_table_with_mor/"; + auto dir = UniqueTestDirectory::Create(); + std::string table_path = dir->Str(); + ASSERT_TRUE(TestUtil::CopyDirectory(test_data_path, table_path)); + auto file_system = std::make_shared(); + + CleanContextBuilder preparation_context_builder(table_path); + ASSERT_OK_AND_ASSIGN(std::unique_ptr preparation_context, + preparation_context_builder.AddOption(Options::FILE_SYSTEM, "local") + .WithOlderThanMs(std::numeric_limits::max()) + .Finish()); + ASSERT_OK_AND_ASSIGN(auto preparation_cleaner, + OrphanFilesCleaner::Create(std::move(preparation_context))); + auto* cleaner_impl = dynamic_cast(preparation_cleaner.get()); + ASSERT_NE(cleaner_impl, nullptr); + + const std::string snapshot_path = PathUtil::JoinPath(table_path, "snapshot/snapshot-3"); + ASSERT_OK_AND_ASSIGN(Snapshot snapshot, Snapshot::FromPath(file_system, snapshot_path)); + std::vector manifest_metas; + ASSERT_OK(cleaner_impl->manifest_list_->ReadDataManifests(snapshot, &manifest_metas)); + ASSERT_FALSE(manifest_metas.empty()); + std::vector manifest_entries; + ASSERT_OK(cleaner_impl->manifest_file_->Read(manifest_metas.front().FileName(), + /*filter=*/nullptr, &manifest_entries)); + ASSERT_FALSE(manifest_entries.empty()); + + const std::string extra_file_name = "data-live-extra.orc"; + const ManifestEntry& source_entry = manifest_entries.front(); + auto file_with_extra = source_entry.File()->CopyWithExtraFiles({extra_file_name}); + ManifestEntry entry_with_extra(source_entry.Kind(), source_entry.Partition(), + source_entry.Bucket(), source_entry.TotalBuckets(), + file_with_extra); + ASSERT_OK_AND_ASSIGN(std::vector extra_manifest_metas, + cleaner_impl->manifest_file_->Write({entry_with_extra})); + using ManifestListWithSize = std::pair; + ASSERT_OK_AND_ASSIGN(ManifestListWithSize extra_manifest_list, + cleaner_impl->manifest_list_->Write(extra_manifest_metas)); + + Snapshot snapshot_with_extra = + WithDataManifests(snapshot, extra_manifest_list.first, extra_manifest_list.second, + snapshot.DeltaManifestList(), snapshot.DeltaManifestListSize()); + ASSERT_OK_AND_ASSIGN(std::string snapshot_json, snapshot_with_extra.ToJsonString()); + ASSERT_OK(file_system->WriteFile(snapshot_path, snapshot_json, /*overwrite=*/true)); + + const std::string extra_file_path = + PathUtil::JoinPath(table_path, "p0=0/p1=0/bucket-0/" + extra_file_name); + ASSERT_OK(file_system->WriteFile(extra_file_path, "extra", /*overwrite=*/false)); + CleanContextBuilder clean_context_builder(table_path); ASSERT_OK_AND_ASSIGN(std::unique_ptr clean_context, - clean_context_builder.AddOption(Options::FILE_SYSTEM, "local").Finish()); - ASSERT_NOK_WITH_MSG(OrphanFilesCleaner::Create(std::move(clean_context)), - "orphan files cleaner only support append table"); + clean_context_builder.AddOption(Options::FILE_SYSTEM, "local") + .WithOlderThanMs(std::numeric_limits::max()) + .Finish()); + ASSERT_OK_AND_ASSIGN(auto cleaner, OrphanFilesCleaner::Create(std::move(clean_context))); + ASSERT_OK_AND_ASSIGN(std::set cleaned_paths, cleaner->Clean()); + + ASSERT_OK_AND_ASSIGN(bool extra_file_exists, file_system->Exists(extra_file_path)); + ASSERT_TRUE(extra_file_exists); + for (const std::string& cleaned_path : cleaned_paths) { + ASSERT_NE(PathUtil::GetName(cleaned_path), extra_file_name); + } } TEST(OrphanFilesCleanerTest, TestTableWithTag) { @@ -245,35 +388,119 @@ TEST(OrphanFilesCleanerTest, TestTableWithChangelog) { } TEST(OrphanFilesCleanerTest, TestTableWithIndexManifest) { - std::string test_data_path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09/"; + std::string test_data_path = + paimon::test::GetDataDir() + "/orc/pk_table_with_mor.db/pk_table_with_mor/"; auto dir = UniqueTestDirectory::Create(); std::string table_path = dir->Str(); ASSERT_TRUE(TestUtil::CopyDirectory(test_data_path, table_path)); auto file_system = std::make_shared(); - auto snapshot_str = R"({ - "version" : 3, - "id" : 6, - "schemaId" : 0, - "baseManifestList" : "manifest-list-f2d59cb8-3ec6-4860-b34b-050b1a533416-0", - "deltaManifestList" : "manifest-list-f2d59cb8-3ec6-4860-b34b-050b1a533416-1", - "changelogManifestList" : null, - "indexManifest" : "index-manifest-bd43150e-cce1-4231-bfc1-8fdc2b0b5994-0", - "commitUser" : "febb1e71-79fc-4abc-9b9d-464ecbc198f7", - "commitIdentifier" : 9223372036854775807, - "commitKind" : "APPEND", - "timeMillis" : 1721615035363, - "totalRecordCount" : 11, - "deltaRecordCount" : 1, - "changelogRecordCount" : 0 -})"; - ASSERT_OK(file_system->WriteFile(PathUtil::JoinPath(table_path, "snapshot/snapshot-6"), - snapshot_str, true)); + auto external_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(external_dir); + std::string external_index_root = "file:" + external_dir->Str(); + + SchemaManager schema_manager(file_system, table_path); + ASSERT_OK_AND_ASSIGN(std::optional> optional_schema, + schema_manager.Latest()); + ASSERT_TRUE(optional_schema); + const std::shared_ptr& table_schema = optional_schema.value(); + std::map raw_options = table_schema->Options(); + raw_options[Options::FILE_SYSTEM] = "local"; + raw_options[Options::GLOBAL_INDEX_EXTERNAL_PATH] = external_index_root; + raw_options[Options::INDEX_FILE_IN_DATA_FILE_DIR] = "true"; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(raw_options)); + std::shared_ptr memory_pool = GetDefaultPool(); + auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); + ASSERT_OK_AND_ASSIGN(std::vector external_paths, options.CreateExternalPaths()); + ASSERT_OK_AND_ASSIGN(std::optional global_index_external_path, + options.CreateGlobalIndexExternalPath()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr path_factory, + FileStorePathFactory::Create( + table_path, arrow_schema, table_schema->PartitionKeys(), + options.GetPartitionDefaultName(), options.GetFileFormat()->Identifier(), + options.DataFilePrefix(), options.LegacyPartitionNameEnabled(), external_paths, + global_index_external_path, options.IndexFileInDataFileDir(), memory_pool)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr index_manifest_file, + IndexManifestFile::Create(file_system, options.GetManifestFormat(), + options.GetManifestCompression(), path_factory, + options.GetBucket(), memory_pool, options)); + ASSERT_OK_AND_ASSIGN(BinaryRow partition, + path_factory->ToBinaryRow({{"p0", "0"}, {"p1", "0"}})); + + const std::string live_internal_name = "btree-global-index-live.index"; + const std::string live_external_name = "btree-global-index-live-external.index"; + const std::string orphan_internal_name = "btree-global-index-orphan.index"; + const std::string live_external_path = + PathUtil::JoinPath(external_index_root, live_external_name); + ASSERT_OK_AND_ASSIGN(IndexManifestEntry live_internal, + CreateSourceBackedBTreeEntry(live_internal_name, partition, + /*external_path=*/std::nullopt, memory_pool)); + ASSERT_OK_AND_ASSIGN(IndexManifestEntry live_external, + CreateSourceBackedBTreeEntry(live_external_name, partition, + live_external_path, memory_pool)); + ASSERT_OK_AND_ASSIGN(IndexManifestEntry orphan_internal, + CreateSourceBackedBTreeEntry(orphan_internal_name, partition, + /*external_path=*/std::nullopt, memory_pool)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr index_path_factory, + path_factory->CreateIndexFileFactory(partition, /*bucket=*/0)); + const std::string live_internal_path = index_path_factory->ToPath(live_internal.index_file); + ASSERT_EQ(PathUtil::JoinPath(table_path, + "p0=0/p1=0/bucket-0/" + live_internal.index_file->FileName()), + live_internal_path); + ASSERT_EQ(live_external_path, index_path_factory->ToPath(live_external.index_file)); + const std::string orphan_internal_path = index_path_factory->ToPath(orphan_internal.index_file); + ASSERT_OK(file_system->WriteFile(live_internal_path, "payload", /*overwrite=*/false)); + ASSERT_OK(file_system->WriteFile(live_external_path, "payload", /*overwrite=*/false)); + ASSERT_OK(file_system->WriteFile(orphan_internal_path, "payload", /*overwrite=*/false)); + + ASSERT_OK_AND_ASSIGN( + std::optional live_manifest, + index_manifest_file->WriteIndexFiles(std::nullopt, {live_internal, live_external})); + ASSERT_TRUE(live_manifest); + ASSERT_OK_AND_ASSIGN(std::optional orphan_manifest, + index_manifest_file->WriteIndexFiles(std::nullopt, {orphan_internal})); + ASSERT_TRUE(orphan_manifest); + + const std::string snapshot_path = PathUtil::JoinPath(table_path, "snapshot/snapshot-3"); + ASSERT_OK_AND_ASSIGN(Snapshot snapshot, Snapshot::FromPath(file_system, snapshot_path)); + Snapshot indexed_snapshot = WithIndexManifest(snapshot, live_manifest.value()); + ASSERT_OK_AND_ASSIGN(std::string indexed_snapshot_json, indexed_snapshot.ToJsonString()); + ASSERT_OK(file_system->WriteFile(snapshot_path, indexed_snapshot_json, /*overwrite=*/true)); CleanContextBuilder clean_context_builder(table_path); ASSERT_OK_AND_ASSIGN(std::unique_ptr clean_context, - clean_context_builder.AddOption(Options::FILE_SYSTEM, "local").Finish()); + clean_context_builder.AddOption(Options::FILE_SYSTEM, "local") + .AddOption(Options::GLOBAL_INDEX_EXTERNAL_PATH, external_index_root) + .AddOption(Options::INDEX_FILE_IN_DATA_FILE_DIR, "true") + .WithOlderThanMs(std::numeric_limits::max()) + .Finish()); ASSERT_OK_AND_ASSIGN(auto cleaner, OrphanFilesCleaner::Create(std::move(clean_context))); - ASSERT_NOK_WITH_MSG(cleaner->Clean(), "OrphanFilesCleaner do not support clean index manifest"); + ASSERT_OK_AND_ASSIGN(std::set cleaned_paths, cleaner->Clean()); + + ASSERT_OK_AND_ASSIGN(bool live_internal_exists, file_system->Exists(live_internal_path)); + ASSERT_TRUE(live_internal_exists); + ASSERT_OK_AND_ASSIGN(bool live_external_exists, file_system->Exists(live_external_path)); + ASSERT_TRUE(live_external_exists); + ASSERT_OK_AND_ASSIGN( + bool live_manifest_exists, + file_system->Exists(path_factory->ToManifestFilePath(live_manifest.value()))); + ASSERT_TRUE(live_manifest_exists); + ASSERT_OK_AND_ASSIGN(bool orphan_internal_exists, file_system->Exists(orphan_internal_path)); + ASSERT_FALSE(orphan_internal_exists); + ASSERT_OK_AND_ASSIGN( + bool orphan_manifest_exists, + file_system->Exists(path_factory->ToManifestFilePath(orphan_manifest.value()))); + ASSERT_FALSE(orphan_manifest_exists); + + std::set cleaned_names; + for (const std::string& path : cleaned_paths) { + cleaned_names.insert(PathUtil::GetName(path)); + } + ASSERT_TRUE(cleaned_names.count(orphan_internal_name)); + ASSERT_TRUE(cleaned_names.count(orphan_manifest.value())); + ASSERT_FALSE(cleaned_names.count(live_internal_name)); + ASSERT_FALSE(cleaned_names.count(live_external_name)); + ASSERT_FALSE(cleaned_names.count(live_manifest.value())); } } // namespace paimon::test diff --git a/src/paimon/core/operation/restore_files.h b/src/paimon/core/operation/restore_files.h index 44c067058..1860b6af9 100644 --- a/src/paimon/core/operation/restore_files.h +++ b/src/paimon/core/operation/restore_files.h @@ -38,12 +38,14 @@ class RestoreFiles { const std::optional& total_buckets, const std::vector>& data_files, const std::shared_ptr& dynamic_bucket_index, - const std::vector>& delete_vectors_index) + const std::vector>& delete_vectors_index, + const std::vector>& primary_key_index_payloads = {}) : snapshot_(snapshot), total_buckets_(total_buckets), data_files_(data_files), dynamic_bucket_index_(dynamic_bucket_index), - delete_vectors_index_(delete_vectors_index) {} + delete_vectors_index_(delete_vectors_index), + primary_key_index_payloads_(primary_key_index_payloads) {} std::optional GetSnapshot() const { return snapshot_; @@ -60,6 +62,9 @@ class RestoreFiles { std::vector> DeleteVectorsIndex() const { return delete_vectors_index_; } + std::vector> PrimaryKeyIndexPayloads() const { + return primary_key_index_payloads_; + } static std::shared_ptr Empty() { return std::make_shared(); @@ -71,6 +76,7 @@ class RestoreFiles { std::vector> data_files_; std::shared_ptr dynamic_bucket_index_; std::vector> delete_vectors_index_; + std::vector> primary_key_index_payloads_; }; } // namespace paimon diff --git a/src/paimon/core/operation/write_restore.h b/src/paimon/core/operation/write_restore.h index a58bedfc8..282129785 100644 --- a/src/paimon/core/operation/write_restore.h +++ b/src/paimon/core/operation/write_restore.h @@ -43,7 +43,8 @@ class WriteRestore { virtual Result LatestCommittedIdentifier(const std::string& user) const = 0; virtual Result> GetRestoreFiles( - const BinaryRow& partition, int32_t bucket, bool scan_delete_vectors_index) const = 0; + const BinaryRow& partition, int32_t bucket, bool scan_delete_vectors_index, + bool scan_source_index_payloads) const = 0; }; } // namespace paimon diff --git a/src/paimon/core/postpone/postpone_bucket_file_store_write.h b/src/paimon/core/postpone/postpone_bucket_file_store_write.h index 1ca65d289..fc993016c 100644 --- a/src/paimon/core/postpone/postpone_bucket_file_store_write.h +++ b/src/paimon/core/postpone/postpone_bucket_file_store_write.h @@ -111,7 +111,8 @@ class PostponeBucketFileStoreWrite : public AbstractFileStoreWrite { : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, table_schema, schema, /*write_schema=*/schema, partition_schema, dv_maintainer_factory, - io_manager, options, ignore_previous_files, is_streaming_mode, + /*primary_key_index_maintainer_factory=*/nullptr, io_manager, + options, ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, executor, pool) {} Result> CreateWriter( diff --git a/src/paimon/core/schema/schema_validation.cpp b/src/paimon/core/schema/schema_validation.cpp index 90f508b47..e6a159999 100644 --- a/src/paimon/core/schema/schema_validation.cpp +++ b/src/paimon/core/schema/schema_validation.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -36,14 +37,19 @@ #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/global_index/btree/btree_defs.h" +#include "paimon/common/options/memory_size.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/vector_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/object_utils.h" +#include "paimon/common/utils/options_utils.h" #include "paimon/common/utils/preconditions.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definition.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" #include "paimon/core/options/changelog_producer.h" #include "paimon/core/options/expire_config.h" #include "paimon/core/options/map_storage_layout.h" @@ -58,6 +64,8 @@ namespace paimon { namespace { +constexpr char kDeletionVectorsMergeOnRead[] = "deletion-vectors.merge-on-read"; + bool ContainsBlobField(const std::shared_ptr& field) { if (BlobUtils::IsBlobField(field)) { return true; @@ -127,6 +135,58 @@ Status ValidatePerLevelOption( return Status::OK(); } +std::vector PrimaryKeyBTreeIndexColumns( + const std::map& options) { + auto iter = options.find(Options::PK_BTREE_INDEX_COLUMNS); + if (iter == options.end()) { + return {}; + } + std::vector columns = StringUtils::Split(iter->second, ",", false); + for (std::string& column : columns) { + StringUtils::Trim(&column); + } + return columns; +} + +bool IsSupportedBTreeIndexType(const std::shared_ptr& type) { + switch (type->id()) { + case arrow::Type::BOOL: + case arrow::Type::INT8: + case arrow::Type::INT16: + case arrow::Type::INT32: + case arrow::Type::INT64: + case arrow::Type::FLOAT: + case arrow::Type::DOUBLE: + case arrow::Type::STRING: + case arrow::Type::DATE32: + case arrow::Type::TIMESTAMP: + case arrow::Type::DECIMAL128: + return true; + default: + return false; + } +} + +Status ValidateBTreeIndexerOptions(const std::map& options) { + PAIMON_ASSIGN_OR_RAISE(std::string cache_size, OptionsUtils::GetValueFromMap( + options, BtreeDefs::kBtreeIndexCacheSize, + BtreeDefs::kDefaultBtreeIndexCacheSize)); + Result parsed_cache_size = MemorySize::ParseBytes(cache_size); + if (!parsed_cache_size.ok()) { + return parsed_cache_size.status().WithMessage(fmt::format( + "Invalid BTree cache size '{}': {}", cache_size, parsed_cache_size.status().message())); + } + PAIMON_ASSIGN_OR_RAISE( + double high_priority_pool_ratio, + OptionsUtils::GetValueFromMap(options, BtreeDefs::kBtreeIndexHighPriorityPoolRatio, + BtreeDefs::kDefaultBtreeIndexHighPriorityPoolRatio)); + if (!std::isfinite(high_priority_pool_ratio) || high_priority_pool_ratio < 0.0 || + high_priority_pool_ratio >= 1.0) { + return Status::Invalid("The BTree high priority pool ratio should be in the range [0, 1)."); + } + return Status::OK(); +} + } // namespace bool SchemaValidation::IsComplexType(const std::shared_ptr& field) { @@ -189,6 +249,7 @@ Status SchemaValidation::ValidateTableSchema(const TableSchema& schema) { if (options.DeletionVectorsEnabled()) { PAIMON_RETURN_NOT_OK(ValidateForDeletionVectors(options)); } + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyBTreeIndexes(schema, options)); PAIMON_RETURN_NOT_OK(ValidateRowTracking(schema, options)); PAIMON_RETURN_NOT_OK(ValidateBlobFields(schema, options)); @@ -368,6 +429,66 @@ Status SchemaValidation::ValidateForDeletionVectors(const CoreOptions& options) "no deletion of old data in this merge engine."); } +Status SchemaValidation::ValidatePrimaryKeyBTreeIndexes(const TableSchema& schema, + const CoreOptions& options) { + std::vector index_columns = PrimaryKeyBTreeIndexColumns(schema.Options()); + if (index_columns.empty()) { + return Status::OK(); + } + + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(schema)); + if (!options.DeletionVectorsEnabled()) { + return Status::Invalid( + "Primary-key BTree indexes require deletion-vectors.enabled = true."); + } + if (schema.PrimaryKeys().empty()) { + return Status::Invalid("Primary-key BTree indexes require a primary-key table."); + } + if (options.GetBucket() <= 0 && !IsPostponeBucketTable(schema, options.GetBucket())) { + return Status::Invalid( + fmt::format("Primary-key BTree indexes require fixed or postpone bucket mode " + "(bucket > 0 or bucket = -2), but bucket is {}.", + options.GetBucket())); + } + PAIMON_ASSIGN_OR_RAISE( + bool deletion_vectors_merge_on_read, + OptionsUtils::GetValueFromMap(schema.Options(), kDeletionVectorsMergeOnRead, false)); + if (deletion_vectors_merge_on_read) { + return Status::Invalid( + "Primary-key BTree indexes require deletion-vectors.merge-on-read = false."); + } + PAIMON_ASSIGN_OR_RAISE(bool pk_clustering_override, + OptionsUtils::GetValueFromMap( + schema.Options(), Options::PK_CLUSTERING_OVERRIDE, false)); + if (pk_clustering_override) { + return Status::Invalid( + "pk-clustering-override is currently unsupported by the C++ commit path, including " + "tables with primary-key BTree indexes."); + } + + for (const std::string& column : index_columns) { + auto field_iter = + std::find_if(schema.Fields().begin(), schema.Fields().end(), + [&column](const DataField& field) { return field.Name() == column; }); + if (field_iter == schema.Fields().end()) { + return Status::Invalid(fmt::format("{} entry '{}' must reference an existing column.", + Options::PK_BTREE_INDEX_COLUMNS, column)); + } + if (!IsSupportedBTreeIndexType(field_iter->Type())) { + return Status::Invalid(fmt::format("{} entry '{}' has unsupported type {}.", + Options::PK_BTREE_INDEX_COLUMNS, column, + field_iter->Type()->ToString())); + } + } + for (const PrimaryKeyIndexDefinition& definition : definitions.Definitions()) { + if (definition.GetFamily() == PrimaryKeyIndexDefinition::Family::BTREE) { + PAIMON_RETURN_NOT_OK(ValidateBTreeIndexerOptions(definition.Options())); + } + } + return Status::OK(); +} + Status SchemaValidation::ValidateSequenceGroup(const TableSchema& schema, const CoreOptions& options) { std::unordered_map> fields2_group; diff --git a/src/paimon/core/schema/schema_validation.h b/src/paimon/core/schema/schema_validation.h index 388ab42e2..e7c17e5a6 100644 --- a/src/paimon/core/schema/schema_validation.h +++ b/src/paimon/core/schema/schema_validation.h @@ -68,6 +68,8 @@ class SchemaValidation { static Status ValidateSequenceGroup(const TableSchema& schema, const CoreOptions& options); static Status ValidateChangelogProducer(const TableSchema& schema, const CoreOptions& options); static Status ValidateForDeletionVectors(const CoreOptions& options); + static Status ValidatePrimaryKeyBTreeIndexes(const TableSchema& schema, + const CoreOptions& options); static Status ValidateRowTracking(const TableSchema& table_schema, const CoreOptions& options); diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 2137970a3..2267669ab 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -30,6 +30,20 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +Result> MakePrimaryKeyBTreeSchema( + std::map options, + const std::vector& primary_keys = {"id"}, + const std::shared_ptr& value_type = arrow::int32()) { + options.emplace(Options::PK_BTREE_INDEX_COLUMNS, "value"); + auto schema = arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false), + arrow::field("value", value_type)}); + return TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, primary_keys, + options); +} + +} // namespace TEST(SchemaValidationTest, TestSimple) { auto f0 = arrow::field("f0", arrow::utf8()); @@ -749,6 +763,102 @@ TEST(SchemaValidationTest, ValidateDeletionVector) { } } +TEST(SchemaValidationTest, ValidatePrimaryKeyBTreeIndexes) { + { + std::map options = { + {Options::BUCKET, "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {"fields.value.pk-btree.index.options", R"({"cache-size":"16 mb"})"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*schema)); + } + { + std::map options = {{Options::BUCKET, "-2"}, + {Options::DELETION_VECTORS_ENABLED, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*schema)); + } + { + std::map options = {{Options::BUCKET, "1"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "require deletion-vectors.enabled = true"); + } + { + std::map options = {{Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::DELETION_VECTORS_ENABLED, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options, /*primary_keys=*/{})); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "require a primary-key table"); + } + { + std::map options = {{Options::BUCKET, "-1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "require fixed or postpone bucket mode"); + } + { + std::map options = {{Options::BUCKET, "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {"deletion-vectors.merge-on-read", "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "require deletion-vectors.merge-on-read = false"); + } + { + std::map options = {{Options::BUCKET, "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {Options::PK_CLUSTERING_OVERRIDE, "true"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "unsupported by the C++ commit path"); + } +} + +TEST(SchemaValidationTest, ValidatePrimaryKeyBTreeIndexColumnsAndOptions) { + std::map valid_options = { + {Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}; + { + std::map options = valid_options; + options.emplace(Options::PK_BTREE_INDEX_COLUMNS, "missing"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "entry 'missing' must reference an existing column"); + } + { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(valid_options, /*primary_keys=*/{"id"}, arrow::binary())); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "entry 'value' has unsupported type binary"); + } + { + std::map options = valid_options; + options.emplace(Options::PK_BTREE_INDEX_COLUMNS, "value,value"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), + "contains duplicate column 'value'"); + } + { + std::map options = valid_options; + options.emplace("fields.value.pk-btree.index.options", R"({"cache-size":"not-a-size"})"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema, + MakePrimaryKeyBTreeSchema(options)); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*schema), "not-a-size"); + } +} + TEST(SchemaValidationTest, ValidateSequenceField) { auto f0 = arrow::field("f0", arrow::utf8()); auto f1 = arrow::field("f1", arrow::int32());