From 3a19320c7b756d97d8794ac5730c47644fc8cddf Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:34:33 +0800 Subject: [PATCH 01/24] feat(realtime): add primary-key in-memory writes Add typed primary-key store creation, an in-memory PK store, and a no-spill writer that materializes sealed mutations through MergeTreeWriter. Keep writer-local compaction disabled, preserve sequence progress across sequential writer handoff, and reject unsupported V1 table options. --- .../realtime/arrow_realtime_store_factory.h | 7 +- include/paimon/realtime/realtime_store.h | 47 +- src/paimon/CMakeLists.txt | 5 + .../core/operation/file_store_write.cpp | 25 +- .../operation/key_value_file_store_write.cpp | 77 ++- .../operation/key_value_file_store_write.h | 8 + .../key_value_file_store_write_test.cpp | 48 ++ .../realtime/arrow_realtime_store_factory.cpp | 55 +- .../realtime/primary_key_realtime_options.cpp | 58 ++ .../realtime/primary_key_realtime_options.h | 31 + .../primary_key_realtime_options_test.cpp | 56 ++ .../realtime/primary_key_realtime_store.cpp | 563 ++++++++++++++++++ .../realtime/primary_key_realtime_store.h | 84 +++ .../primary_key_realtime_store_test.cpp | 244 ++++++++ .../realtime/realtime_append_only_writer.cpp | 11 +- .../core/realtime/realtime_context_impl.cpp | 67 ++- .../core/realtime/realtime_context_impl.h | 18 +- .../core/realtime/realtime_context_test.cpp | 126 +--- .../realtime/realtime_primary_key_writer.cpp | 249 ++++++++ .../realtime/realtime_primary_key_writer.h | 89 +++ 20 files changed, 1690 insertions(+), 178 deletions(-) create mode 100644 src/paimon/core/realtime/primary_key_realtime_options.cpp create mode 100644 src/paimon/core/realtime/primary_key_realtime_options.h create mode 100644 src/paimon/core/realtime/primary_key_realtime_options_test.cpp create mode 100644 src/paimon/core/realtime/primary_key_realtime_store.cpp create mode 100644 src/paimon/core/realtime/primary_key_realtime_store.h create mode 100644 src/paimon/core/realtime/primary_key_realtime_store_test.cpp create mode 100644 src/paimon/core/realtime/realtime_primary_key_writer.cpp create mode 100644 src/paimon/core/realtime/realtime_primary_key_writer.h diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 4d65743ab..da1b8de36 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -26,11 +26,8 @@ namespace paimon { /// Factory for Paimon's default Arrow-backed `RealtimeStore`. class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: - /// Creates an Arrow-backed store for one partition and bucket. - Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) override; + /// Creates the built-in append or in-memory primary-key store. + Result> Create(RealtimeStoreCreateRequest&& request) override; }; } // namespace paimon diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index d02952acd..1e53c173e 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -25,8 +25,10 @@ #include #include #include +#include #include +#include "arrow/c/abi.h" #include "paimon/reader/batch_reader.h" #include "paimon/realtime/offset_range.h" #include "paimon/record_batch.h" @@ -41,6 +43,31 @@ namespace paimon { class MemoryPool; class Predicate; +struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig { + StatisticsMode statistics_mode; +}; + +struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig { + std::vector primary_keys; + /// Largest sequence restored from the committed snapshot. A PK store assigns one contiguous + /// sequence to every mutation in `Write` order, starting at the next value, and rejects + /// `Write` before the assigned sequence would exceed `INT64_MAX - 1`. + int64_t restore_max_sequence_number; +}; + +using RealtimeStoreCreateConfig = + std::variant; + +struct PAIMON_EXPORT RealtimeStoreCreateRequest { + /// Complete table write schema whose ownership is transferred to the factory. + std::unique_ptr<::ArrowSchema> write_schema; + std::map options; + std::shared_ptr memory_pool; + std::map partition; + int32_t bucket = -1; + RealtimeStoreCreateConfig mode_config; +}; + /// A table record batch and its framework-assigned contiguous offset range. /// /// The batch contains only table write fields. Row `i` is associated with @@ -133,8 +160,11 @@ class PAIMON_EXPORT RealtimeStore { /// /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by /// `context.read_schema` except a duplicate `_VALUE_KIND`. Concatenating all returned readers - /// must produce every matching row once. Paimon retains `view` for the lifetime of the - /// resulting framework reader. + /// must produce every matching row once. Primary-key readers additionally provide a non-null + /// `_SEQUENCE_NUMBER` when requested, are individually sorted by primary key, and contain at + /// most one mutation per key. Assigned sequences remain stable across views and queries; + /// readers need not be globally sorted with one another. Paimon retains `view` for the lifetime + /// of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; @@ -157,16 +187,9 @@ class PAIMON_EXPORT RealtimeStoreFactory { public: virtual ~RealtimeStoreFactory() = default; - /// Creates a store configured with the supplied schema, statistics, options, and memory pool. - /// @param write_schema Complete table write schema whose ownership is transferred to the - /// factory. The factory may consume it or retain it in the created store. - /// @param statistics_mode Framework-parsed statistics collection mode. - /// @param options Effective table options available to the store. - /// @param memory_pool Memory pool provided by the write context. - virtual Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) = 0; + /// Creates a store for the requested table mode and partition-bucket. + /// The factory consumes `request`, including ownership of `request.write_schema`. + virtual Result> Create(RealtimeStoreCreateRequest&& request) = 0; }; } // namespace paimon diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index a9810424a..69deab92f 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -378,9 +378,12 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/primary_key_realtime_store.cpp + core/realtime/primary_key_realtime_options.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp + core/realtime/realtime_primary_key_writer.cpp core/postpone/postpone_bucket_writer.cpp core/schema/arrow_schema_validator.cpp core/schema/schema_manager.cpp @@ -780,6 +783,8 @@ if(PAIMON_BUILD_TESTS) core/manifest/index_manifest_file_handler_test.cpp core/memory/writer_memory_manager_test.cpp core/realtime/arrow_realtime_store_test.cpp + core/realtime/primary_key_realtime_store_test.cpp + core/realtime/primary_key_realtime_options_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 6807ae35e..fb83c254c 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -36,6 +36,7 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" +#include "paimon/core/realtime/primary_key_realtime_options.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" @@ -197,7 +198,26 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - return Status::Invalid("real-time write currently supports append tables only"); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options)); + if (ignore_previous_files) { + return Status::NotImplemented( + "PK realtime v1 requires restore from the latest snapshot"); + } + if (!ctx->GetWriteSchema().empty()) { + return Status::NotImplemented( + "PK realtime v1 does not support a custom write schema"); + } + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager->LatestSnapshot()); + if (latest_snapshot) { + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap realtime_committed_offsets, + RealtimeCommitProperties::ReadOffsets( + latest_snapshot, options.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(ctx->GetRealtimeContext())); + PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( + latest_snapshot->Id(), realtime_committed_offsets)); + } } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( @@ -253,7 +273,8 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRootPath(), 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->GetExecutor(), ctx->GetMemoryPool()); + ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), + ctx->GetMemoryPool()); } } 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 08c5ea0c3..4456ee1c2 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -18,21 +18,29 @@ #include "paimon/core/operation/key_value_file_store_write.h" +#include +#include #include +#include "arrow/c/bridge.h" #include "paimon/common/data/binary_row.h" +#include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/mergetree/levels.h" #include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/realtime/realtime_context.h" namespace arrow { class Schema; @@ -60,6 +68,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, 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, @@ -67,6 +76,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( 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), user_defined_seq_comparator_(user_defined_seq_comparator), merge_function_wrapper_(merge_function_wrapper), @@ -74,7 +84,25 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( options_, key_comparator_, user_defined_seq_comparator_, compaction_metrics_, table_schema_, schema_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, root_path_, pool_)), - logger_(Logger::GetLogger("KeyValueFileStoreWrite")) {} + logger_(Logger::GetLogger("KeyValueFileStoreWrite")) { + if (realtime_context_) { + writer_memory_manager_ = std::make_unique(); + } +} + +Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + if (!realtime_context_) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), + options_.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context_)); + return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets); +} Result> KeyValueFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { @@ -109,19 +137,48 @@ Result> KeyValueFileStoreWrite::CreateWriter( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr levels, Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); - auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr compact_manager, - compact_manager_factory_->CreateCompactManager(partition, bucket, compact_strategy, - compact_executor_, levels, dv_maintainer)); + std::map partition_map; + int64_t materialized_max_sequence_number = restore_max_seq_number; + if (realtime_context_) { + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + file_store_path_factory_->GeneratePartitionVector(partition)); + partition_map = + std::map(partition_values.begin(), partition_values.end()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context_)); + const RealtimePartitionBucket partition_bucket(partition_map, bucket); + materialized_max_sequence_number = realtime_context_impl->GetMaterializedMaxSequenceNumber( + partition_bucket, restore_max_seq_number); + if (materialized_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK sequence number has reached INT64_MAX"); + } + } + std::shared_ptr compact_manager; + if (realtime_context_) { + compact_manager = std::make_shared(); + } else { + auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); + PAIMON_ASSIGN_OR_RAISE(compact_manager, compact_manager_factory_->CreateCompactManager( + partition, bucket, compact_strategy, + compact_executor_, levels, dv_maintainer)); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( - restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, - user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, - options_, compact_manager, io_manager_, enable_multi_thread_spill_, pool_)); - return writer; + materialized_max_sequence_number, trimmed_primary_keys, data_file_path_factory, + key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, + table_schema_->Id(), schema_, options_, compact_manager, + realtime_context_ ? nullptr : io_manager_, enable_multi_thread_spill_, pool_)); + if (!realtime_context_) { + return std::shared_ptr(std::move(writer)); + } + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, c_write_schema.get())); + return RealtimePrimaryKeyWriter::Create( + partition_map, bucket, std::move(c_write_schema), trimmed_primary_keys, realtime_context_, + writer, options_.ToMap(), pool_, materialized_max_sequence_number); } Status KeyValueFileStoreWrite::Close() { 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 14457590f..66c362f2e 100644 --- a/src/paimon/core/operation/key_value_file_store_write.h +++ b/src/paimon/core/operation/key_value_file_store_write.h @@ -45,6 +45,7 @@ class SnapshotManager; class SchemaManager; class TableSchema; class IOManager; +class RealtimeContext; struct KeyValue; template class MergeFunctionWrapper; @@ -65,8 +66,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, 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); + Status RefreshCommittedSnapshot(int64_t snapshot_id) override; Status Close() override; private: @@ -79,8 +82,13 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { Result> CreateFileStoreScan( const std::shared_ptr& filter) const override; + bool IsRealtimeWrite() const override { + return realtime_context_ != nullptr; + } + private: bool enable_multi_thread_spill_; + std::shared_ptr realtime_context_; std::shared_ptr key_comparator_; std::shared_ptr user_defined_seq_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index 35d938af7..45462ea6e 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -53,6 +53,7 @@ #include "paimon/format/reader_builder.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" #include "paimon/testing/utils/test_helper.h" @@ -303,6 +304,53 @@ TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenLookupEnabl ASSERT_EQ(commit_messages.size(), 1); } +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { + const std::map options = { + {Options::BUCKET, "1"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + }; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), /*nullable=*/false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithTempDirectory(dir->Str()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + ASSERT_OK(writer->Write(MakeBatch(schema, R"([ + [1, "old"], + [2, "two"], + [1, "new"] + ])"))); + ASSERT_OK_AND_ASSIGN(std::vector progresses, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progresses.size()); + ASSERT_EQ(OffsetRange(0, 3), progresses[0].offset_range); + std::shared_ptr commit_message = + std::dynamic_pointer_cast(progresses[0].commit_message); + ASSERT_NE(nullptr, commit_message); + int64_t row_count = 0; + for (const std::shared_ptr& file : + commit_message->GetNewFilesIncrement().NewFiles()) { + row_count += file->row_count; + } + ASSERT_EQ(2, row_count); + ASSERT_EQ(0, TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str())); + ASSERT_OK(writer->Close()); +} + TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenDefaultCompactRewriterPathEnabled) { ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 1d7219c41..e6e22edfd 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -21,29 +21,66 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/scope_guard.h" +#include "paimon/core/core_options.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" namespace paimon { Result> ArrowRealtimeStoreFactory::Create( - std::unique_ptr write_schema, StatisticsMode statistics_mode, - const std::map&, const std::shared_ptr& memory_pool) { - if (!write_schema || !write_schema->release) { + RealtimeStoreCreateRequest&& request) { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("real-time store write schema is null"); } - ScopeGuard schema_guard([schema = write_schema.get()]() { ArrowSchemaRelease(schema); }); - if (!memory_pool) { + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + if (!request.memory_pool) { return Status::Invalid("real-time store memory pool is null"); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, - arrow::ImportSchema(write_schema.get())); - std::shared_ptr arrow_pool = GetArrowPool(memory_pool); - return std::make_shared(imported_schema, statistics_mode, memory_pool, - arrow_pool); + arrow::ImportSchema(request.write_schema.get())); + if (std::holds_alternative(request.mode_config)) { + const AppendRealtimeStoreCreateConfig& append_config = + std::get(request.mode_config); + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + return std::make_shared(imported_schema, append_config.statistics_mode, + request.memory_pool, arrow_pool); + } + + const PrimaryKeyRealtimeStoreCreateConfig& primary_key_config = + std::get(request.mode_config); + std::vector key_fields; + key_fields.reserve(primary_key_config.primary_keys.size()); + for (const std::string& primary_key : primary_key_config.primary_keys) { + const int32_t field_index = imported_schema->GetFieldIndex(primary_key); + if (field_index < 0) { + return Status::Invalid("primary key ", primary_key, " is missing from write schema"); + } + key_fields.emplace_back(field_index, imported_schema->field(field_index)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); + auto merge_function_wrapper_factory = []() { + auto merge_function = std::make_unique( + /*ignore_delete=*/false); + return std::make_shared(std::move(merge_function)); + }; + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(request.options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(imported_schema, primary_key_config.primary_keys, + key_comparator, merge_function_wrapper_factory, + primary_key_config.restore_max_sequence_number, + core_options.GetReadBatchSize(), request.memory_pool)); + return std::shared_ptr(std::move(store)); } } // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_options.cpp b/src/paimon/core/realtime/primary_key_realtime_options.cpp new file mode 100644 index 000000000..e9779a59e --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_options.cpp @@ -0,0 +1,58 @@ +/* + * 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/realtime/primary_key_realtime_options.h" + +#include "paimon/core/core_options.h" + +namespace paimon { + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime v1 requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime v1 does not support data evolution"); + } + if (!options.GetFieldsSequenceGroups().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence groups"); + } + if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || + options.AggregationRemoveRecordOnDelete() || + !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { + return Status::NotImplemented("PK realtime v1 requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime v1 supports only ascending sequence.field.sort-order"); + } + if (options.NeedLookup() || options.DeletionVectorsEnabled() || + options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); + } + return Status::OK(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_options.h b/src/paimon/core/realtime/primary_key_realtime_options.h new file mode 100644 index 000000000..a16d35778 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_options.h @@ -0,0 +1,31 @@ +/* + * 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 "paimon/status.h" + +namespace paimon { + +class CoreOptions; + +/// Validates the table options supported by the in-memory PK realtime V1 path. +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options); + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_options_test.cpp b/src/paimon/core/realtime/primary_key_realtime_options_test.cpp new file mode 100644 index 000000000..5d3ea7f67 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_options_test.cpp @@ -0,0 +1,56 @@ +/* + * 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/realtime/primary_key_realtime_options.h" + +#include +#include +#include + +#include "paimon/core/core_options.h" +#include "paimon/defs.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(PrimaryKeyRealtimeOptionsTest, TestSupportedOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options)); +} + +TEST(PrimaryKeyRealtimeOptionsTest, TestUnsupportedOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options)); + } +} + +} // namespace paimon::test diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 000000000..84afb97a4 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,563 @@ +/* + * 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/realtime/primary_key_realtime_store.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/data/binary_row_writer.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.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/core/io/key_value_in_memory_record_reader.h" +#include "paimon/core/io/key_value_projection_consumer.h" +#include "paimon/core/io/key_value_projection_reader.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h" +#include "paimon/macros.h" + +namespace paimon { +namespace { + +uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t result = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + result += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + result += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + result += GetArrayMemoryUsage(data->dictionary); + } + return result; +} + +struct StoredBatch { + std::shared_ptr data; + std::vector row_kinds; + OffsetRange offset_range; + int64_t first_sequence_number; + uint64_t memory_usage; +}; +using BatchGroup = std::vector>; + +class Segment final : public RealtimeSegmentHandle { + public: + Segment(const OffsetRange& offset_range, + std::vector>&& batches) + : offset_range_(offset_range), batches_(std::move(batches)) {} + + OffsetRange GetOffsetRange() const override { + return offset_range_; + } + + const std::vector>& Batches() const { + return batches_; + } + + uint64_t GetMemoryUsage() const { + uint64_t result = 0; + for (const std::shared_ptr& batch : batches_) { + result += batch->memory_usage; + } + return result; + } + + private: + OffsetRange offset_range_; + std::vector> batches_; +}; + +class PrimaryKeyRealtimeReadView final : public RealtimeReadView { + public: + explicit PrimaryKeyRealtimeReadView(std::vector&& groups) + : groups_(std::move(groups)) { + if (!groups_.empty()) { + offset_range_ = OffsetRange(groups_.front().front()->offset_range.begin, + groups_.back().back()->offset_range.end); + } + } + + std::optional GetOffsetRange() const override { + return offset_range_; + } + + const std::vector& Groups() const { + return groups_; + } + + private: + std::vector groups_; + std::optional offset_range_; +}; + +class CommitBatchReader final : public BatchReader { + public: + CommitBatchReader(const std::shared_ptr& segment, + const std::shared_ptr& arrow_pool) + : segment_(segment), arrow_pool_(arrow_pool), metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!segment_ || next_batch_ >= static_cast(segment_->Batches().size())) { + return MakeEofBatch(); + } + const std::shared_ptr& stored = segment_->Batches()[next_batch_++]; + const int64_t row_count = stored->data->length(); + arrow::Int8Builder row_kind_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Reserve(row_count)); + if (stored->row_kinds.empty()) { + for (int64_t i = 0; i < row_count; ++i) { + row_kind_builder.UnsafeAppend(static_cast(RecordBatch::RowKind::INSERT)); + } + } else { + for (RecordBatch::RowKind row_kind : stored->row_kinds) { + row_kind_builder.UnsafeAppend(static_cast(row_kind)); + } + } + std::shared_ptr row_kind_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); + arrow::ArrayVector arrays = {std::move(row_kind_array)}; + arrays.insert(arrays.end(), stored->data->fields().begin(), stored->data->fields().end()); + arrow::FieldVector fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; + const arrow::FieldVector& value_fields = stored->data->struct_type()->fields(); + fields.insert(fields.end(), value_fields.begin(), value_fields.end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr output, + arrow::StructArray::Make(arrays, fields)); + auto c_array = std::make_unique(); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*output, c_array.get(), c_schema.get())); + return ReadBatch(std::move(c_array), std::move(c_schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + + void Close() override { + segment_.reset(); + } + + private: + std::shared_ptr segment_; + std::shared_ptr arrow_pool_; + std::shared_ptr metrics_; + int32_t next_batch_ = 0; +}; + +class KeyRangeBatchReader final : public BatchReader, public PrimaryKeyRangeProvider { + public: + KeyRangeBatchReader(std::unique_ptr&& reader, + const std::shared_ptr& min_key, + const std::shared_ptr& max_key) + : reader_(std::move(reader)), min_key_(min_key), max_key_(max_key) {} + + Result NextBatch() override { + return reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + std::shared_ptr GetMinKey() const override { + return min_key_; + } + + std::shared_ptr GetMaxKey() const override { + return max_key_; + } + + private: + std::unique_ptr reader_; + std::shared_ptr min_key_; + std::shared_ptr max_key_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + Impl(const std::shared_ptr& write_schema, std::vector primary_keys, + const std::shared_ptr& key_comparator, + const std::function>()>& + merge_function_wrapper_factory, + int64_t next_sequence_number, int32_t read_batch_size, + const std::shared_ptr& memory_pool) + : write_schema_(write_schema), + primary_keys_(std::move(primary_keys)), + key_comparator_(key_comparator), + merge_function_wrapper_factory_(merge_function_wrapper_factory), + next_sequence_number_(next_sequence_number), + read_batch_size_(read_batch_size), + memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)) {} + + Result> CopyKey(const InternalRow& key) const { + auto result = std::make_shared(static_cast(primary_keys_.size())); + BinaryRowWriter writer(result.get(), /*initial_size=*/128, memory_pool_.get()); + writer.Reset(); + for (int32_t index = 0; index < static_cast(primary_keys_.size()); ++index) { + std::shared_ptr field = + write_schema_->GetFieldByName(primary_keys_[index]); + PAIMON_ASSIGN_OR_RAISE(InternalRow::FieldGetterFunc getter, + InternalRow::CreateFieldGetter(index, field->type(), + /*use_view=*/true)); + PAIMON_ASSIGN_OR_RAISE(BinaryRowWriter::FieldSetterFunc setter, + BinaryRowWriter::CreateFieldSetter(index, field->type())); + setter(getter(key), &writer); + } + writer.Complete(); + return std::static_pointer_cast(result); + } + + Result, std::shared_ptr>> GetKeyRange( + const std::shared_ptr& values) const { + arrow::ArrayVector key_arrays; + key_arrays.reserve(primary_keys_.size()); + for (const std::string& primary_key : primary_keys_) { + std::shared_ptr key_array = values->GetFieldByName(primary_key); + if (!key_array) { + return Status::Invalid("primary key is missing from PK query batch: ", primary_key); + } + key_arrays.push_back(std::move(key_array)); + } + auto context = std::make_shared(key_arrays, memory_pool_); + int64_t min_row = 0; + int64_t max_row = 0; + for (int64_t row = 1; row < values->length(); ++row) { + ColumnarRowRef current(context, row); + ColumnarRowRef min_key(context, min_row); + ColumnarRowRef max_key(context, max_row); + if (key_comparator_->CompareTo(current, min_key) < 0) { + min_row = row; + } + if (key_comparator_->CompareTo(current, max_key) > 0) { + max_row = row; + } + } + ColumnarRowRef min_key(context, min_row); + ColumnarRowRef max_key(context, max_row); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr copied_min, CopyKey(min_key)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr copied_max, CopyKey(max_key)); + return std::make_pair(std::move(copied_min), std::move(copied_max)); + } + + Status Write(RealtimeWriteBatch&& write_batch) { + if (!write_batch.batch || !write_batch.batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t row_count = write_batch.batch->GetData()->length; + if (row_count <= 0 || write_batch.offset_range.begin < 0 || + write_batch.offset_range.Count() != row_count) { + return Status::Invalid("PK real-time offset range does not match batch row count"); + } + const std::vector& row_kinds = write_batch.batch->GetRowKind(); + if (!row_kinds.empty() && static_cast(row_kinds.size()) != row_count) { + return Status::Invalid("PK real-time row-kind count does not match batch row count"); + } + for (RecordBatch::RowKind row_kind : row_kinds) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, + RowKind::FromByteValue(static_cast(row_kind))); + static_cast(validated); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr imported, + arrow::ImportArray(write_batch.batch->GetData(), + arrow::struct_(write_schema_->fields()))); + if (!imported || imported->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr values = + checked_pointer_cast(imported); + PAIMON_RETURN_NOT_OK_FROM_ARROW(values->ValidateFull()); + + std::lock_guard lock(mutex_); + if (last_offset_ && write_batch.offset_range.begin != last_offset_.value()) { + return Status::Invalid("PK real-time offset ranges must be contiguous"); + } + if (row_count > std::numeric_limits::max() - next_sequence_number_) { + return Status::Invalid("PK sequence range exceeds INT64_MAX"); + } + auto stored = std::make_shared( + StoredBatch{std::move(values), row_kinds, write_batch.offset_range, + next_sequence_number_, GetArrayMemoryUsage(imported->data())}); + building_batches_.push_back(std::move(stored)); + building_memory_usage_ += building_batches_.back()->memory_usage; + last_offset_ = write_batch.offset_range.end; + next_sequence_number_ += row_count; + return Status::OK(); + } + + Result>> SealForCommit() { + std::lock_guard lock(mutex_); + if (building_batches_.empty()) { + return std::optional>(); + } + const OffsetRange range(building_batches_.front()->offset_range.begin, + building_batches_.back()->offset_range.end); + auto segment = std::make_shared(range, std::move(building_batches_)); + sealed_segments_.push_back(segment); + building_batches_.clear(); + building_memory_usage_ = 0; + return std::optional>(std::move(segment)); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) { + std::shared_ptr typed = std::dynamic_pointer_cast(segment); + if (!typed) { + return Status::Invalid("segment was not created by the PK real-time store"); + } + std::vector> result; + result.push_back(std::make_unique(typed, arrow_pool_)); + return result; + } + + Result> AcquireReadView() { + std::lock_guard lock(mutex_); + std::vector groups; + groups.reserve(sealed_segments_.size() + (building_batches_.empty() ? 0 : 1)); + for (const std::shared_ptr& segment : sealed_segments_) { + groups.push_back(segment->Batches()); + } + if (!building_batches_.empty()) { + groups.push_back(building_batches_); + } + return std::shared_ptr(new PrimaryKeyRealtimeReadView(std::move(groups))); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t lower, + const RealtimeQueryContext& context) { + std::shared_ptr typed = + std::dynamic_pointer_cast(view); + if (!typed) { + return Status::Invalid("read view was not created by the PK real-time store"); + } + if (!context.read_schema || !context.read_schema->release) { + return Status::Invalid("PK real-time query read schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr requested, + arrow::ImportSchema(context.read_schema)); + arrow::FieldVector output_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; + std::vector projection = {KeyValueProjectionConsumer::kValueKindProjection}; + for (const std::shared_ptr& field : requested->fields()) { + if (field->name() == SpecialFields::ValueKind().Name()) { + continue; + } + output_fields.push_back(field); + if (field->name() == SpecialFields::SequenceNumber().Name()) { + projection.push_back(KeyValueProjectionConsumer::kSequenceNumberProjection); + continue; + } + const int32_t index = write_schema_->GetFieldIndex(field->name()); + if (index < 0) { + return Status::Invalid("PK real-time query field is missing from write schema: ", + field->name()); + } + projection.push_back(index); + } + + std::vector> result; + for (const BatchGroup& group : typed->Groups()) { + std::vector> batch_readers; + std::shared_ptr min_key; + std::shared_ptr max_key; + for (const std::shared_ptr& batch : group) { + if (batch->offset_range.end <= lower) { + continue; + } + const int64_t offset = std::max(0, lower - batch->offset_range.begin); + const int64_t length = batch->data->length() - offset; + std::shared_ptr sliced = batch->data->Slice(offset, length); + std::shared_ptr selected = + checked_pointer_cast(sliced); + using KeyRange = + std::pair, std::shared_ptr>; + PAIMON_ASSIGN_OR_RAISE(KeyRange key_range, GetKeyRange(selected)); + if (!min_key || key_comparator_->CompareTo(*key_range.first, *min_key) < 0) { + min_key = key_range.first; + } + if (!max_key || key_comparator_->CompareTo(*key_range.second, *max_key) > 0) { + max_key = key_range.second; + } + std::vector selected_kinds; + if (!batch->row_kinds.empty()) { + selected_kinds.assign(batch->row_kinds.begin() + offset, + batch->row_kinds.end()); + } + std::unique_ptr reader = + std::make_unique( + batch->first_sequence_number + offset, selected, selected_kinds, + primary_keys_, /*user_defined_sequence_fields=*/std::vector(), + /*sequence_fields_ascending=*/true, key_comparator_, memory_pool_); + std::shared_ptr> batch_merge = + merge_function_wrapper_factory_(); + if (!batch_merge) { + return Status::Invalid("merge function wrapper factory returned null"); + } + batch_readers.push_back(std::make_unique( + std::move(reader), key_comparator_, batch_merge)); + } + if (batch_readers.empty()) { + continue; + } + std::shared_ptr> group_merge = + merge_function_wrapper_factory_(); + if (!group_merge) { + return Status::Invalid("merge function wrapper factory returned null"); + } + auto merged = std::make_unique( + std::move(batch_readers), key_comparator_, + /*user_defined_seq_comparator=*/nullptr, group_merge); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr projected, + KeyValueProjectionReader::Create(std::move(merged), arrow::schema(output_fields), + projection, read_batch_size_, memory_pool_)); + result.push_back( + std::make_unique(std::move(projected), min_key, max_key)); + } + return result; + } + + Status AdvanceCommittedOffset(int64_t committed_end_offset) { + std::lock_guard lock(mutex_); + sealed_segments_.erase( + std::remove_if(sealed_segments_.begin(), sealed_segments_.end(), + [committed_end_offset](const std::shared_ptr& segment) { + return segment->GetOffsetRange().end <= committed_end_offset; + }), + sealed_segments_.end()); + return Status::OK(); + } + + uint64_t GetMemoryUsage() const { + std::lock_guard lock(mutex_); + uint64_t result = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_segments_) { + result += segment->GetMemoryUsage(); + } + return result; + } + + private: + std::shared_ptr write_schema_; + std::vector primary_keys_; + std::shared_ptr key_comparator_; + std::function>()> + merge_function_wrapper_factory_; + int64_t next_sequence_number_; + int32_t read_batch_size_; + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + mutable std::mutex mutex_; + std::vector> building_batches_; + std::vector> sealed_segments_; + uint64_t building_memory_usage_ = 0; + std::optional last_offset_; +}; + +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& write_schema, + const std::vector& primary_keys, + const std::shared_ptr& key_comparator, + const std::function>()>& + merge_function_wrapper_factory, + int64_t restore_max_sequence_number, int32_t read_batch_size, + const std::shared_ptr& memory_pool) { + if (!write_schema || primary_keys.empty() || !key_comparator || + !merge_function_wrapper_factory || !memory_pool || read_batch_size <= 0) { + return Status::Invalid("PK real-time store requires schema, keys, merge helpers, and pool"); + } + if (restore_max_sequence_number < -1) { + return Status::Invalid("PK restore max sequence number must be at least -1"); + } + if (restore_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK sequence number has reached INT64_MAX"); + } + for (const std::string& key : primary_keys) { + if (write_schema->GetFieldIndex(key) < 0) { + return Status::Invalid("primary key ", key, " is missing from write schema"); + } + } + auto impl = std::make_unique( + write_schema, primary_keys, key_comparator, merge_function_wrapper_factory, + restore_max_sequence_number + 1, read_batch_size, memory_pool); + return std::shared_ptr(new PrimaryKeyRealtimeStore(std::move(impl))); +} + +PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) + : impl_(std::move(impl)) {} + +PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; + +Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { + return impl_->Write(std::move(batch)); +} + +Result>> +PrimaryKeyRealtimeStore::SealForCommit() { + return impl_->SealForCommit(); +} + +Result>> PrimaryKeyRealtimeStore::CreateCommitReaders( + const std::shared_ptr& segment) { + return impl_->CreateCommitReaders(segment); +} + +Result> PrimaryKeyRealtimeStore::AcquireReadView() { + return impl_->AcquireReadView(); +} + +Result>> PrimaryKeyRealtimeStore::CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) { + return impl_->CreateQueryReaders(view, offset_begin, context); +} + +Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_offset) { + return impl_->AdvanceCommittedOffset(committed_offset); +} + +uint64_t PrimaryKeyRealtimeStore::GetMemoryUsage() const { + return impl_->GetMemoryUsage(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h new file mode 100644 index 000000000..05225ed19 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,84 @@ +/* + * 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/realtime/realtime_store.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class FieldsComparator; +struct KeyValue; +class MemoryPool; +class InternalRow; +template +class MergeFunctionWrapper; + +/// Optional metadata exposed by PK query readers with a known inclusive key range. +class PrimaryKeyRangeProvider { + public: + virtual ~PrimaryKeyRangeProvider() = default; + + virtual std::shared_ptr GetMinKey() const = 0; + virtual std::shared_ptr GetMaxKey() const = 0; +}; + +/// In-memory store for primary-key real-time writes. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& write_schema, + const std::vector& primary_keys, + const std::shared_ptr& key_comparator, + const std::function>()>& + merge_function_wrapper_factory, + int64_t restore_max_sequence_number, int32_t read_batch_size, + const std::shared_ptr& memory_pool); + + ~PrimaryKeyRealtimeStore() override; + + Status Write(RealtimeWriteBatch&& batch) override; + Result>> SealForCommit() override; + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override; + Status AdvanceCommittedOffset(int64_t committed_offset) override; + uint64_t GetMemoryUsage() const override; + + private: + class Impl; + explicit PrimaryKeyRealtimeStore(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp new file mode 100644 index 000000000..9da272e0f --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,244 @@ +/* + * 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/realtime/primary_key_realtime_store.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +class PrimaryKeyRealtimeStoreTest : public testing::Test { + public: + void SetUp() override { + pool_ = std::shared_ptr(GetMemoryPool()); + schema_ = arrow::schema( + {arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(key_comparator_, + FieldsComparator::Create({DataField(0, schema_->field(0))}, + /*is_ascending_order=*/true)); + auto merge_factory = []() { + auto merge_function = + std::make_unique(/*ignore_delete=*/false); + return std::make_shared(std::move(merge_function)); + }; + ASSERT_OK_AND_ASSIGN( + store_, PrimaryKeyRealtimeStore::Create(schema_, {"id"}, key_comparator_, merge_factory, + /*restore_max_sequence_number=*/4, + /*read_batch_size=*/1024, pool_)); + } + + std::unique_ptr MakeBatch( + const std::string& json, const std::vector& row_kinds = {}) const { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json) + .ValueOrDie(); + ArrowArray c_array; + EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + RecordBatchBuilder builder(&c_array); + builder.SetRowKinds(row_kinds); + return builder.Finish().value(); + } + + std::unique_ptr MakeReadSchema(bool include_sequence) const { + arrow::FieldVector fields; + if (include_sequence) { + fields.push_back( + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())); + } + fields.insert(fields.end(), schema_->fields().begin(), schema_->fields().end()); + auto c_schema = std::make_unique(); + EXPECT_TRUE(arrow::ExportSchema(*arrow::schema(fields), c_schema.get()).ok()); + return c_schema; + } + + void AssertReaderOutput(BatchReader* reader, const std::shared_ptr& type, + const std::string& json) const { + ASSERT_NE(nullptr, reader); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + arrow::Result> imported_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported_result.ok()) << imported_result.status().ToString(); + std::shared_ptr actual = std::move(imported_result).ValueOrDie(); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(type, json).ValueOrDie(); + ASSERT_TRUE(actual->Equals(*expected)) + << "expected: " << expected->ToString() << ", actual: " << actual->ToString(); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + reader->Close(); + } + + std::shared_ptr CommitType() const { + return arrow::struct_({ + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), + schema_->field(0), + schema_->field(1), + }); + } + + std::shared_ptr QueryType() const { + return arrow::struct_({ + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()), + schema_->field(0), + schema_->field(1), + }); + } + + protected: + std::shared_ptr pool_; + std::shared_ptr schema_; + std::shared_ptr key_comparator_; + std::shared_ptr store_; +}; + +TEST_F(PrimaryKeyRealtimeStoreTest, TestWriteAndSeal) { + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store_->SealForCommit()); + ASSERT_FALSE(segment.has_value()); + ASSERT_NOK_WITH_MSG(store_->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), + "write batch is null"); + ASSERT_NOK_WITH_MSG( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "a"], [2, "b"]])"), OffsetRange(0, 0)}), + "offset range does not match batch row count"); + + ASSERT_OK( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "a"], [2, "b"]])"), OffsetRange(0, 2)})); + ASSERT_NOK_WITH_MSG( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[4, "d"]])"), OffsetRange(3, 4)}), + "offset ranges must be contiguous"); + ASSERT_OK(store_->Write(RealtimeWriteBatch{MakeBatch(R"([[3, "c"]])"), OffsetRange(2, 3)})); + + ASSERT_OK_AND_ASSIGN(segment, store_->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); + ASSERT_GT(store_->GetMemoryUsage(), 0); + + auto merge_factory = []() { + auto merge_function = std::make_unique(/*ignore_delete=*/false); + return std::make_shared(std::move(merge_function)); + }; + ASSERT_NOK_WITH_MSG(PrimaryKeyRealtimeStore::Create( + schema_, {"id"}, key_comparator_, merge_factory, + /*restore_max_sequence_number=*/-2, /*read_batch_size=*/1024, pool_), + "restore max sequence number must be at least -1"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestCommitReaderPreservesMutations) { + ASSERT_OK(store_->Write( + RealtimeWriteBatch{MakeBatch(R"([[2, "old"], [1, "one"], [2, "new"]])", + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER}), + OffsetRange(0, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store_->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + AssertReaderOutput(readers[0].get(), CommitType(), + R"([[0, 2, "old"], [0, 1, "one"], [2, 2, "new"]])"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestMutationMerge) { + ASSERT_OK(store_->Write( + RealtimeWriteBatch{MakeBatch(R"([[1, "old"], [2, "two"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store_->Write(RealtimeWriteBatch{ + MakeBatch(R"([[1, "new"], [2, "gone"]])", + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE}), + OffsetRange(2, 4)})); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + AssertReaderOutput(readers[0].get(), QueryType(), R"([[2, 7, 1, "new"], [3, 8, 2, "gone"]])"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestReadViewLifecycle) { + ASSERT_OK(store_->Write(RealtimeWriteBatch{MakeBatch(R"([[10, "a"], [11, "b"], [12, "c"]])"), + OffsetRange(10, 13)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store_->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + ASSERT_EQ(std::optional(OffsetRange(10, 13)), view->GetOffsetRange()); + + ASSERT_OK(store_->AdvanceCommittedOffset(13)); + ASSERT_EQ(0, store_->GetMemoryUsage()); + ASSERT_OK( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[13, "later"]])"), OffsetRange(13, 14)})); + + std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, /*offset_begin=*/11, context)); + ASSERT_EQ(1, readers.size()); + AssertReaderOutput(readers[0].get(), QueryType(), R"([[0, 6, 11, "b"], [0, 7, 12, "c"]])"); + + std::unique_ptr empty_schema = MakeReadSchema(/*include_sequence=*/true); + context.read_schema = empty_schema.get(); + ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/13, context)); + ASSERT_TRUE(readers.empty()); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryKeyRange) { + ASSERT_OK(store_->Write( + RealtimeWriteBatch{MakeBatch(R"([[5, "five"], [1, "one"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store_->SealForCommit()); + ASSERT_OK(store_->Write( + RealtimeWriteBatch{MakeBatch(R"([[9, "nine"], [7, "seven"]])"), OffsetRange(2, 4)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(2, readers.size()); + auto* first_range = dynamic_cast(readers[0].get()); + auto* second_range = dynamic_cast(readers[1].get()); + ASSERT_NE(nullptr, first_range); + ASSERT_NE(nullptr, second_range); + ASSERT_EQ(1, first_range->GetMinKey()->GetLong(0)); + ASSERT_EQ(5, first_range->GetMaxKey()->GetLong(0)); + ASSERT_EQ(7, second_range->GetMinKey()->GetLong(0)); + ASSERT_EQ(9, second_range->GetMaxKey()->GetLong(0)); +} + +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 9d519d791..21d6cfb74 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -47,7 +47,7 @@ Result> RealtimeAppendOnlyWriter::Crea std::unique_ptr<::ArrowSchema> write_schema, const std::shared_ptr& realtime_context, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, StatisticsMode statistics_mode, + const std::shared_ptr& input_schema, const std::map& options, const std::shared_ptr& memory_pool) { if (!realtime_context) { @@ -55,10 +55,11 @@ Result> RealtimeAppendOnlyWriter::Crea } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, RealtimeContextImpl::Cast(realtime_context)); - PAIMON_ASSIGN_OR_RAISE( - RealtimeStoreState store_state, - realtime_context_impl->GetOrCreateRealtimeStore(partition, bucket, std::move(write_schema), - statistics_mode, options, memory_pool)); + RealtimeStoreCreateRequest request{ + std::move(write_schema), options, memory_pool, partition, bucket, + AppendRealtimeStoreCreateConfig{statistics_mode}}; + PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore(std::move(request))); return std::shared_ptr(new RealtimeAppendOnlyWriter( store_state.store, file_writer, input_schema, store_state.initial_offset, memory_pool)); } diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index f6bad5cf1..0a367b2cd 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -78,19 +78,16 @@ Status RealtimeContextImpl::Start() { } Result RealtimeContextImpl::GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) { + RealtimeStoreCreateRequest&& request) { std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); - const RealtimePartitionBucket key(partition, bucket); + const RealtimePartitionBucket key(request.partition, request.bucket); int64_t initial_offset = 0; auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); + if (request.write_schema) { + ArrowSchemaRelease(request.write_schema.get()); } return Status::Invalid("real-time offset has reached INT64_MAX"); } @@ -98,8 +95,8 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( } auto iter = stores_.find(key); if (iter != stores_.end()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); + if (request.write_schema) { + ArrowSchemaRelease(request.write_schema.get()); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, iter->second->AcquireReadView()); @@ -119,9 +116,8 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( } return RealtimeStoreState{iter->second, initial_offset}; } - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr store, - factory_->Create(std::move(write_schema), statistics_mode, options, memory_pool)); + Result> store_result = factory_->Create(std::move(request)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); stores_.emplace(key, store); if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); @@ -129,6 +125,27 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( return RealtimeStoreState{std::move(store), initial_offset}; } +int64_t RealtimeContextImpl::GetMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t restored_max_sequence_number) { + std::lock_guard lock(mutex_); + auto [iter, inserted] = + materialized_max_sequence_numbers_.emplace(partition_bucket, restored_max_sequence_number); + if (!inserted && restored_max_sequence_number > iter->second) { + iter->second = restored_max_sequence_number; + } + return iter->second; +} + +void RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + std::lock_guard lock(mutex_); + auto [iter, inserted] = + materialized_max_sequence_numbers_.emplace(partition_bucket, max_sequence_number); + if (!inserted && max_sequence_number > iter->second) { + iter->second = max_sequence_number; + } +} + Result> RealtimeContextImpl::AcquireReadViews() { std::lock_guard lock(mutex_); std::vector result; @@ -230,28 +247,12 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, if (partition_bucket.bucket < 0 || committed_end_offset < 0) { return Status::Invalid("invalid partition-bucket committed offset"); } - } - // Only stores created by this context can contain state which cannot be restored in - // place. Offsets for other partition-buckets are reference state for lazy store creation - // and may be removed or rolled back without rebuilding the context. - std::lock_guard registry_lock(mutex_); - for (const auto& store_entry : stores_) { - const RealtimePartitionBucket& partition_bucket = store_entry.first; auto previous_iter = committed_offsets_.find(partition_bucket); - if (previous_iter == committed_offsets_.end()) { - continue; - } - - auto current_iter = committed_offsets.find(partition_bucket); - if (current_iter == committed_offsets.end()) { - return Status::Invalid( - "real-time committed progress removed an active partition-bucket; recreate " - "RealtimeContext"); - } - if (current_iter->second < previous_iter->second) { - return Status::Invalid( - "real-time committed offset moved backwards for an active partition-bucket; " - "recreate RealtimeContext"); + if (previous_iter != committed_offsets_.end()) { + if (committed_end_offset < previous_iter->second) { + return Status::Invalid( + "real-time partition-bucket committed offset cannot move backwards"); + } } } committed_offsets_ = committed_offsets; diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 66c324cab..45d07deeb 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -32,8 +32,8 @@ #include #include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" #include "paimon/result.h" -#include "paimon/statistics_mode.h" #include "paimon/visibility.h" struct ArrowSchema; @@ -65,11 +65,13 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { static Result> Cast( const std::shared_ptr& context); - Result GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool); + Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); + + int64_t GetMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, + int64_t restored_max_sequence_number); + + void AdvanceMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, + int64_t max_sequence_number); Result> AcquireReadViews(); @@ -79,9 +81,6 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Status ReleaseReadView(const std::string& opaque_ticket); - // Returns an error requiring a new context if a newer snapshot removes or moves committed - // progress backwards for a store created by this context. Progress for inactive stores is - // only reference state and can be replaced in place. Status AdvanceCommittedProgress(int64_t snapshot_id, const RealtimeOffsetMap& committed_offsets); @@ -103,6 +102,7 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::mutex mutex_; std::mutex progress_mutex_; std::map> stores_; + std::map materialized_max_sequence_numbers_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; // Progress already reflected in stores owned by this context. diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 017820fd4..33701afac 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -91,14 +91,11 @@ class TestingRealtimeStore : public RealtimeStore { class TestingRealtimeStoreFactory : public RealtimeStoreFactory { public: - Result> Create(std::unique_ptr write_schema, - StatisticsMode, - const std::map&, - const std::shared_ptr&) override { - if (!write_schema || !write_schema->release) { + Result> Create(RealtimeStoreCreateRequest&& request) override { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("testing write schema is null"); } - ArrowSchemaRelease(write_schema.get()); + ArrowSchemaRelease(request.write_schema.get()); auto store = std::make_shared(); stores.push_back(store); return store; @@ -122,20 +119,28 @@ Result> CreateContext( return RealtimeContextImpl::Cast(context); } -TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { +Result GetOrCreateAppendStore( + const std::shared_ptr& context, + const std::map& partition, int32_t bucket, + std::unique_ptr write_schema, const std::map& options, + const std::shared_ptr& memory_pool) { + return context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(write_schema), options, memory_pool, partition, bucket, + AppendRealtimeStoreCreateConfig{StatisticsMode::NONE}}); +} + +TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); std::shared_ptr pool = GetDefaultPool(); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {{"k", "v"}}, pool)); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState first_state, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, + MakeWriteSchema(), {{"k", "v"}}, pool)); ASSERT_EQ(0, first_state.initial_offset); ASSERT_OK_AND_ASSIGN( RealtimeStoreState first_again_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, pool)); ASSERT_EQ(first_state.store, first_again_state.store); ASSERT_EQ(0, first_again_state.initial_offset); ASSERT_EQ(1, factory->stores.size()); @@ -143,12 +148,10 @@ TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { ASSERT_OK_AND_ASSIGN( RealtimeStoreState second_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 1, MakeWriteSchema(), {}, pool)); ASSERT_OK_AND_ASSIGN( RealtimeStoreState third_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-03"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, {{"dt", "2026-08-03"}}, 0, MakeWriteSchema(), {}, pool)); ASSERT_NE(first_state.store, second_state.store); ASSERT_NE(first_state.store, third_state.store); ASSERT_EQ(3, factory->stores.size()); @@ -171,10 +174,8 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); ASSERT_EQ(2, factory->stores.size()); ASSERT_NOK_WITH_MSG(context->AdvanceCommittedProgress(-1, {}), @@ -193,8 +194,7 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_OK_AND_ASSIGN( RealtimeStoreState restored_state, - context->GetOrCreateRealtimeStore({{"dt", "unknown"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, {{"dt", "unknown"}}, 0, MakeWriteSchema(), {}, pool)); ASSERT_EQ(9, restored_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress( @@ -211,41 +211,15 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } -TEST(RealtimeContextTest, TestRemovedInactivePartitionDoesNotRequireReopen) { - auto factory = std::make_shared(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - const std::map active_partition = {{"dt", "2026-08-02"}}; - const std::map inactive_partition = {{"dt", "2026-08-03"}}; - const RealtimePartitionBucket active_partition_bucket(active_partition, /*bucket=*/0); - const RealtimePartitionBucket inactive_partition_bucket(inactive_partition, /*bucket=*/0); - - ASSERT_OK(context->AdvanceCommittedProgress( - 5, {{active_partition_bucket, /*offset=*/7}, {inactive_partition_bucket, /*offset=*/9}})); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState active_state, - context->GetOrCreateRealtimeStore(active_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_EQ(7, active_state.initial_offset); - - ASSERT_OK(context->AdvanceCommittedProgress(6, {{active_partition_bucket, /*offset=*/7}})); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState inactive_state, - context->GetOrCreateRealtimeStore(inactive_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_EQ(0, inactive_state.initial_offset); -} - TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 2, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, partition, 2, MakeWriteSchema(), {}, pool)); ASSERT_EQ(3, factory->stores.size()); factory->stores[1]->fail_next_advance = true; @@ -259,9 +233,9 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_TRUE(factory->stores[1]->committed_offsets.empty()); ASSERT_EQ(std::vector({9}), factory->stores[2]->committed_offsets); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState failed_store_state, - context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState failed_store_state, + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); ASSERT_EQ(8, failed_store_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress(5, committed_offsets)); @@ -271,45 +245,11 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } -TEST(RealtimeContextTest, TestRequiresReopenWhenCommittedProgressMovesBackwards) { - auto factory = std::make_shared(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - const std::map first_partition = {{"dt", "2026-08-02"}}; - const std::map second_partition = {{"dt", "2026-08-03"}}; - const RealtimePartitionBucket first_partition_bucket(first_partition, /*bucket=*/0); - const RealtimePartitionBucket second_partition_bucket(second_partition, /*bucket=*/0); - - ASSERT_OK(context->GetOrCreateRealtimeStore(first_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(second_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->AdvanceCommittedProgress( - 5, {{first_partition_bucket, /*offset=*/7}, {second_partition_bucket, /*offset=*/9}})); - ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); - - ASSERT_NOK_WITH_MSG( - context->AdvanceCommittedProgress( - 6, {{first_partition_bucket, /*offset=*/6}, {second_partition_bucket, /*offset=*/10}}), - "recreate RealtimeContext"); - ASSERT_NOK_WITH_MSG( - context->AdvanceCommittedProgress(6, {{first_partition_bucket, /*offset=*/10}}), - "recreate RealtimeContext"); - ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); - - ASSERT_OK(context->AdvanceCommittedProgress( - 6, {{first_partition_bucket, /*offset=*/10}, {second_partition_bucket, /*offset=*/11}})); - ASSERT_EQ(std::vector({7, 10}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9, 11}), factory->stores[1]->committed_offsets); -} - TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); @@ -332,8 +272,8 @@ TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { TEST(RealtimeContextTest, TestExpiresAbandonedReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp new file mode 100644 index 000000000..2ebcede82 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,249 @@ +/* + * 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/realtime/realtime_primary_key_writer.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.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/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" +#include "paimon/realtime/realtime_context.h" + +namespace paimon { + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + std::unique_ptr<::ArrowSchema> write_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& realtime_context, + const std::shared_ptr& merge_tree_writer, + const std::map& options, + const std::shared_ptr& memory_pool, int64_t restore_max_sequence_number) { + ScopeGuard schema_guard([schema = write_schema.get()]() { + if (schema && schema->release) { + ArrowSchemaRelease(schema); + } + }); + if (!realtime_context) { + return Status::Invalid("PK real-time context is null"); + } + if (!merge_tree_writer) { + return Status::Invalid("PK real-time merge-tree writer is null"); + } + if (!write_schema || !write_schema->release) { + return Status::Invalid("PK real-time write schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, + arrow::ImportSchema(write_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*imported_schema, write_schema.get())); + RealtimeStoreCreateRequest request{ + std::move(write_schema), + options, + memory_pool, + partition, + bucket, + PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys, restore_max_sequence_number}}; + schema_guard.Release(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore(std::move(request))); + return std::shared_ptr( + new RealtimePrimaryKeyWriter(store_state.store, merge_tree_writer, realtime_context_impl, + RealtimePartitionBucket(partition, bucket), imported_schema, + store_state.initial_offset, memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, int64_t next_offset, + const std::shared_ptr& memory_pool) + : memory_pool_(memory_pool), + realtime_store_(realtime_store), + merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), + write_schema_(write_schema), + next_offset_(next_offset) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t row_count = batch->GetData()->length; + if (row_count == 0) { + return Status::OK(); + } + std::lock_guard lock(realtime_store_mutex_); + if (row_count > std::numeric_limits::max() - next_offset_) { + return Status::Invalid("real-time offset range exceeds INT64_MAX"); + } + const OffsetRange range(next_offset_, next_offset_ + row_count); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{std::move(batch), range})); + next_offset_ += row_count; + return Status::OK(); +} + +Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard lock(prepare_mutex_); + std::optional> segment; + { + std::lock_guard realtime_store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> sealed_segment, + realtime_store_->SealForCommit()); + segment = std::move(sealed_segment); + } + if (segment) { + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value())); + } + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, + merge_tree_writer_->PrepareCommit(wait_compaction)); + if (segment) { + const std::vector>& new_files = + increment.GetNewFilesIncrement().NewFiles(); + if (!new_files.empty()) { + realtime_context_->AdvanceMaterializedMaxSequenceNumber( + partition_bucket_, DataFileMeta::GetMaxSequenceNumber(new_files)); + } + increment.SetRealtimeOffsetRange(segment.value()->GetOffsetRange()); + } + return increment; +} + +Status RealtimePrimaryKeyWriter::FlushSegment( + const std::shared_ptr& segment) { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + realtime_store_->CreateCommitReaders(segment)); + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } + ConcatBatchReader reader(std::move(readers), memory_pool_); + ScopeGuard reader_guard([&reader]() { reader.Close(); }); + const OffsetRange offset_range = segment->GetOffsetRange(); + int64_t emitted_rows = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader.NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!imported || imported->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time store commit reader returned a non-StructArray"); + } + std::shared_ptr struct_array = + checked_pointer_cast(imported); + std::shared_ptr value_kind = + struct_array->GetFieldByName(SpecialFields::ValueKind().Name()); + if (!value_kind || value_kind->type_id() != arrow::Type::INT8) { + return Status::Invalid( + "PK real-time store commit reader must return an INT8 _VALUE_KIND field"); + } + std::shared_ptr encoded_row_kinds = + checked_pointer_cast(value_kind); + std::vector row_kinds; + row_kinds.reserve(static_cast(encoded_row_kinds->length())); + for (int64_t i = 0; i < encoded_row_kinds->length(); ++i) { + if (encoded_row_kinds->IsNull(i)) { + return Status::Invalid("PK real-time store commit reader returned a null row kind"); + } + PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, + RowKind::FromByteValue(encoded_row_kinds->Value(i))); + row_kinds.push_back(static_cast(row_kind->ToByteValue())); + } + PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( + struct_array, SpecialFields::ValueKind().Name())); + if (!struct_array->type()->Equals(arrow::struct_(write_schema_->fields()))) { + return Status::Invalid( + "PK real-time store commit reader schema does not match table write schema"); + } + const int64_t row_count = struct_array->length(); + if (row_count > offset_range.Count() - emitted_rows) { + return Status::Invalid( + "PK real-time store commit readers returned more rows than the sealed offset " + "range"); + } + emitted_rows += row_count; + if (row_count == 0) { + continue; + } + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, output.get())); + RecordBatchBuilder builder(output.get()); + builder.SetRowKinds(row_kinds); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr record_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(merge_tree_writer_->Write(std::move(record_batch))); + } + if (emitted_rows != offset_range.Count()) { + return Status::Invalid( + "PK real-time store commit readers returned fewer rows than the sealed offset range"); + } + return Status::OK(); +} + +Status RealtimePrimaryKeyWriter::Compact(bool) { + return Status::Invalid("PK real-time write does not support explicit compaction"); +} + +uint64_t RealtimePrimaryKeyWriter::GetMemoryUsage() const { + return realtime_store_->GetMemoryUsage(); +} + +Status RealtimePrimaryKeyWriter::FlushMemory() { + return Status::OK(); +} + +Result RealtimePrimaryKeyWriter::CompactNotCompleted() { + return merge_tree_writer_->CompactNotCompleted(); +} + +Status RealtimePrimaryKeyWriter::Sync() { + return merge_tree_writer_->Sync(); +} + +Status RealtimePrimaryKeyWriter::Close() { + return merge_tree_writer_->Close(); +} + +std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { + return merge_tree_writer_->GetMetrics(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h new file mode 100644 index 000000000..fa057e079 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" + +struct ArrowSchema; + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; +class MergeTreeWriter; +class RealtimeContext; +class RealtimeContextImpl; + +/// Primary-key real-time writer backed by an in-memory mutation indexer. +class RealtimePrimaryKeyWriter final : public BatchWriter { + public: + static Result> Create( + const std::map& partition, int32_t bucket, + std::unique_ptr<::ArrowSchema> write_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& realtime_context, + const std::shared_ptr& merge_tree_writer, + const std::map& options, + const std::shared_ptr& memory_pool, int64_t restore_max_sequence_number); + + Status Write(std::unique_ptr&& batch) override; + Result PrepareCommit(bool wait_compaction) override; + Status Compact(bool full_compaction) override; + uint64_t GetMemoryUsage() const override; + Status FlushMemory() override; + Result CompactNotCompleted() override; + Status Sync() override; + Status Close() override; + std::shared_ptr GetMetrics() const override; + + private: + RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + int64_t next_offset, const std::shared_ptr& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment); + + std::shared_ptr memory_pool_; + std::shared_ptr realtime_store_; + std::shared_ptr merge_tree_writer_; + std::shared_ptr realtime_context_; + RealtimePartitionBucket partition_bucket_; + std::shared_ptr write_schema_; + int64_t next_offset_; + std::mutex realtime_store_mutex_; + std::mutex prepare_mutex_; +}; + +} // namespace paimon From 789da66ca4741de6240db301ff751b3b4ee30d54 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:34:33 +0800 Subject: [PATCH 02/24] feat(read): merge primary-key realtime memory with snapshots Capture partition-bucket read views in realtime splits and merge PK memory readers with snapshot data by key range. Retain read views for reader lifetime, defer ticket consumption until vector reader construction succeeds, and apply predicates after PK deduplication. --- .../core/operation/merge_file_split_read.cpp | 274 ++++++++++++++++++ .../core/operation/merge_file_split_read.h | 18 ++ .../table/source/key_value_table_read.cpp | 264 +++++++++++++++++ .../core/table/source/key_value_table_read.h | 7 + .../core/table/source/realtime_table_scan.cpp | 2 +- .../core/table/source/realtime_table_scan.h | 2 +- src/paimon/core/table/source/table_scan.cpp | 7 +- 7 files changed, 569 insertions(+), 5 deletions(-) diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index b753ea431..8d8367e39 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -30,6 +30,7 @@ #include "arrow/c/bridge.h" #include "arrow/type.h" #include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/table/special_fields.h" @@ -78,6 +79,273 @@ struct KeyValue; template class MergeFunctionWrapper; +namespace { + +/// Concatenates merge readers whose key ranges are ordered and non-overlapping, preserving one +/// projection pipeline without merging independent disk-only components. +class ConcatNonOverlappingMergeReaders final : public SortMergeReader { + public: + explicit ConcatNonOverlappingMergeReaders( + std::vector>&& readers) + : readers_(std::move(readers)) {} + + Result> NextBatch() override { + while (current_ < readers_.size()) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + readers_[current_]->NextBatch()); + if (iterator) { + return iterator; + } + readers_[current_]->Close(); + ++current_; + } + return std::unique_ptr(); + } + + void Close() override { + while (current_ < readers_.size()) { + readers_[current_++]->Close(); + } + } + + std::shared_ptr GetReaderMetrics() const override { + return MetricsImpl::CollectReadMetrics(readers_); + } + + private: + std::vector> readers_; + size_t current_ = 0; +}; + +} // namespace + +class MergeFileSplitRead::RealtimeReaderBuilder { + public: + static Result> Create( + MergeFileSplitRead* owner, const std::vector>& disk_splits, + std::vector&& additional_readers) { + RealtimeReaderBuilder builder(owner); + if (disk_splits.empty()) { + std::vector> readers; + readers.reserve(additional_readers.size()); + for (AdditionalKeyValueReader& additional : additional_readers) { + readers.push_back(std::move(additional.reader)); + } + return builder.CreateMergedReader(std::move(readers)); + } + + PAIMON_RETURN_NOT_OK(builder.CollectDiskInputs(disk_splits)); + builder.AddRangeInputs(std::move(additional_readers)); + return builder.CreateReader(); + } + + private: + struct RangeInput { + std::shared_ptr min_key; + std::shared_ptr max_key; + std::vector disk_runs; + std::unique_ptr additional_reader; + }; + + explicit RealtimeReaderBuilder(MergeFileSplitRead* owner) : owner_(owner) {} + + Status CollectDiskInputs(const std::vector>& disk_splits) { + first_split_ = std::dynamic_pointer_cast(disk_splits.front()); + if (!first_split_) { + return Status::Invalid("merge input disk split is not a data split"); + } + const BinaryRow& partition = first_split_->Partition(); + const int32_t bucket = first_split_->Bucket(); + PAIMON_ASSIGN_OR_RAISE(data_file_path_factory_, + owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); + + std::vector> data_files; + std::vector> deletion_files; + for (const std::shared_ptr& disk_split : disk_splits) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + if (!data_split || !(data_split->Partition() == partition) || + data_split->Bucket() != bucket) { + return Status::Invalid("merge input disk splits do not share a partition-bucket"); + } + if (!data_split->BeforeFiles().empty() || data_split->IsStreaming() || + data_split->Bucket() == BucketModeDefine::POSTPONE_BUCKET) { + return Status::Invalid("additional merge input requires fixed-bucket batch splits"); + } + const std::vector>& split_files = data_split->DataFiles(); + const std::vector>& split_deletion_files = + data_split->DeletionFiles(); + if (!split_deletion_files.empty() && + split_deletion_files.size() != split_files.size()) { + return Status::Invalid( + "merge input disk split deletion files must be empty or match data files"); + } + data_files.insert(data_files.end(), split_files.begin(), split_files.end()); + if (split_deletion_files.empty()) { + deletion_files.insert(deletion_files.end(), split_files.size(), std::nullopt); + } else { + deletion_files.insert(deletion_files.end(), split_deletion_files.begin(), + split_deletion_files.end()); + } + } + + dv_factory_ = DeletionVector::CreateFactory( + owner_->options_.GetFileSystem(), + DeletionVector::CreateDeletionFileMap(data_files, deletion_files), owner_->pool_); + std::vector> disk_sections = + IntervalPartition(data_files, owner_->key_comparator_).Partition(); + inputs_.reserve(disk_sections.size()); + for (std::vector& section : disk_sections) { + std::shared_ptr min_file = section.front().Files().front(); + std::shared_ptr max_file = min_file; + for (const SortedRun& run : section) { + for (const std::shared_ptr& file : run.Files()) { + if (owner_->key_comparator_->CompareTo(file->min_key, min_file->min_key) < 0) { + min_file = file; + } + if (owner_->key_comparator_->CompareTo(file->max_key, max_file->max_key) > 0) { + max_file = file; + } + } + } + inputs_.push_back(RangeInput{std::shared_ptr(min_file, &min_file->min_key), + std::shared_ptr(max_file, &max_file->max_key), + std::move(section), nullptr}); + } + return Status::OK(); + } + + void AddRangeInputs(std::vector&& additional_readers) { + inputs_.reserve(inputs_.size() + additional_readers.size()); + for (AdditionalKeyValueReader& additional : additional_readers) { + has_unknown_range_ |= !additional.min_key || !additional.max_key; + inputs_.push_back(RangeInput{additional.min_key, additional.max_key, + /*disk_runs=*/{}, std::move(additional.reader)}); + } + } + + Result> CreateDiskReader(const SortedRun& run) { + return owner_->CreateReaderForRun(first_split_->Partition(), run, dv_factory_, + owner_->predicate_for_keys_, data_file_path_factory_); + } + + Result> CreateMergedReader( + std::vector>&& record_readers) { + if (record_readers.empty()) { + return std::make_unique(std::vector>{}, + owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + owner_->CreateSortMergeReader(std::move(record_readers))); + return CreateProjectedReader(std::move(sort_merge_reader)); + } + + Result> CreateProjectedReader( + std::unique_ptr&& sort_merge_reader) { + if (!owner_->force_keep_delete_) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } + + std::unique_ptr projection_reader; + if (!owner_->context_->EnableMultiThreadRowToBatch()) { + PAIMON_ASSIGN_OR_RAISE( + projection_reader, + KeyValueProjectionReader::Create( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), owner_->pool_)); + } else { + const int32_t thread_number = owner_->context_->GetRowToBatchThreadNumber(); + assert(thread_number > 0); + projection_reader = std::make_unique( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), thread_number, owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(projection_reader, + owner_->ApplyPredicateFilterIfNeeded( + std::move(projection_reader), owner_->context_->GetPredicate())); + return std::make_unique(std::move(projection_reader), + owner_->pool_); + } + + Result> CreateUnknownRangeReader() { + std::vector> readers; + for (RangeInput& input : inputs_) { + for (const SortedRun& run : input.disk_runs) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + CreateDiskReader(run)); + readers.push_back(std::move(disk_reader)); + } + if (input.additional_reader) { + readers.push_back(std::move(input.additional_reader)); + } + } + return CreateMergedReader(std::move(readers)); + } + + Result> CreateKnownRangeReader() { + std::sort(inputs_.begin(), inputs_.end(), + [this](const RangeInput& lhs, const RangeInput& rhs) { + return owner_->key_comparator_->CompareTo(*lhs.min_key, *rhs.min_key) < 0; + }); + std::vector> components; + std::shared_ptr component_max_key; + for (RangeInput& input : inputs_) { + if (components.empty() || + owner_->key_comparator_->CompareTo(*input.min_key, *component_max_key) > 0) { + components.emplace_back(); + component_max_key = input.max_key; + } else if (owner_->key_comparator_->CompareTo(*input.max_key, *component_max_key) > 0) { + component_max_key = input.max_key; + } + components.back().push_back(std::move(input)); + } + + std::vector> component_readers; + component_readers.reserve(components.size()); + for (std::vector& component : components) { + if (component.size() == 1 && !component.front().additional_reader) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr disk_component, + owner_->CreateSortMergeReaderForSection( + component.front().disk_runs, first_split_->Partition(), dv_factory_, + component.front().disk_runs.size() == 1 ? owner_->context_->GetPredicate() + : owner_->predicate_for_keys_, + data_file_path_factory_, /*drop_delete=*/false)); + component_readers.push_back(std::move(disk_component)); + continue; + } + + std::vector> readers; + for (RangeInput& input : component) { + for (const SortedRun& run : input.disk_runs) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + CreateDiskReader(run)); + readers.push_back(std::move(disk_reader)); + } + if (input.additional_reader) { + readers.push_back(std::move(input.additional_reader)); + } + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr component_reader, + owner_->CreateSortMergeReader(std::move(readers))); + component_readers.push_back(std::move(component_reader)); + } + return CreateProjectedReader( + std::make_unique(std::move(component_readers))); + } + + Result> CreateReader() { + return has_unknown_range_ ? CreateUnknownRangeReader() : CreateKnownRangeReader(); + } + + MergeFileSplitRead* owner_; + std::shared_ptr first_split_; + std::shared_ptr data_file_path_factory_; + DeletionVector::Factory dv_factory_; + std::vector inputs_; + bool has_unknown_range_ = false; +}; + Result> MergeFileSplitRead::Create( const std::shared_ptr& path_factory, const std::shared_ptr& context, @@ -158,6 +426,12 @@ Result> MergeFileSplitRead::CreateReader( return std::make_unique(std::move(batch_reader), pool_); } +Result> MergeFileSplitRead::CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector&& additional_readers) { + return RealtimeReaderBuilder::Create(this, disk_splits, std::move(additional_readers)); +} + void MergeFileSplitRead::SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper) { merge_function_wrapper_ = merge_function_wrapper; diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index d4bfa727c..8c541ec6f 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -55,6 +55,7 @@ class FieldsComparator; class FileBatchReader; class FileStorePathFactory; class InternalReadContext; +class InternalRow; class MemoryPool; class SchemaManager; class SortedRun; @@ -65,6 +66,12 @@ struct KeyValue; template class MergeFunctionWrapper; +struct AdditionalKeyValueReader { + std::unique_ptr reader; + std::shared_ptr min_key; + std::shared_ptr max_key; +}; + /// If the class name below is enclosed in parentheses, it might be present in the read path; /// otherwise, it must be present in the read path. /// @@ -117,10 +124,21 @@ class MergeFileSplitRead : public AbstractSplitRead { return value_schema_; } + std::shared_ptr GetKeySchema() const { + return key_schema_; + } + + /// Merges ordinary disk splits with generic additional sorted KeyValue readers. + Result> CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector&& additional_readers); + void SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper); private: + class RealtimeReaderBuilder; + Result> CreateMergeReader( const std::shared_ptr& data_split, const std::shared_ptr& data_file_path_factory); diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 208807493..770caf1ca 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -20,12 +20,28 @@ #include "paimon/core/table/source/key_value_table_read.h" #include +#include +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/key_value.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" +#include "paimon/core/table/source/realtime_split.h" #include "paimon/status.h" namespace paimon { @@ -34,6 +50,163 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +struct ColumnarBatchContext; + +namespace { + +class QueryBatchKeyValueReader final : public KeyValueRecordReader { + public: + QueryBatchKeyValueReader(std::unique_ptr&& reader, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& pool) + : reader_(std::move(reader)), + key_schema_(key_schema), + value_schema_(value_schema), + pool_(pool) {} + + Result> NextBatch() override; + std::shared_ptr GetReaderMetrics() const override; + void Close() override; + + private: + class Iterator; + + std::unique_ptr reader_; + std::shared_ptr key_schema_; + std::shared_ptr value_schema_; + std::shared_ptr pool_; + std::shared_ptr values_; + std::shared_ptr sequences_; + std::shared_ptr row_kinds_; + std::shared_ptr key_context_; + std::shared_ptr value_context_; +}; + +class QueryBatchKeyValueReader::Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(QueryBatchKeyValueReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->values_->length(); + } + + Result Next() override { + if (reader_->sequences_->IsNull(cursor_) || reader_->row_kinds_->IsNull(cursor_)) { + return Status::Invalid("PK merge metadata must not be null"); + } + PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kinds_->Value(cursor_))); + const int64_t sequence = reader_->sequences_->Value(cursor_); + std::shared_ptr key = + std::make_shared(reader_->key_context_, cursor_); + auto value = std::make_unique(reader_->value_context_, cursor_++); + return KeyValue(row_kind, sequence, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + QueryBatchKeyValueReader* reader_; + int64_t cursor_ = 0; +}; + +Result> QueryBatchKeyValueReader::NextBatch() { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return std::unique_ptr(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr input = + std::dynamic_pointer_cast(imported); + if (!input) { + return Status::Invalid("PK merge input is not a StructArray"); + } + sequences_ = std::dynamic_pointer_cast( + input->GetFieldByName(SpecialFields::SequenceNumber().Name())); + row_kinds_ = std::dynamic_pointer_cast( + input->GetFieldByName(SpecialFields::ValueKind().Name())); + if (!sequences_ || !row_kinds_) { + return Status::Invalid("PK merge input is missing sequence or value-kind metadata"); + } + PAIMON_ASSIGN_OR_RAISE(input, ArrowUtils::RemoveFieldFromStructArray( + input, SpecialFields::SequenceNumber().Name())); + PAIMON_ASSIGN_OR_RAISE( + values_, ArrowUtils::RemoveFieldFromStructArray(input, SpecialFields::ValueKind().Name())); + if (!ArrowUtils::EqualsIgnoreNullable(values_->type(), + arrow::struct_(value_schema_->fields()))) { + return Status::Invalid("PK merge input value schema does not match the table read schema"); + } + arrow::ArrayVector key_fields; + key_fields.reserve(key_schema_->num_fields()); + for (const std::shared_ptr& field : key_schema_->fields()) { + std::shared_ptr key = values_->GetFieldByName(field->name()); + if (!key) { + return Status::Invalid("PK merge input is missing key field ", field->name()); + } + key_fields.push_back(std::move(key)); + } + key_context_ = std::make_shared(key_fields, pool_); + value_context_ = std::make_shared(values_->fields(), pool_); + return std::make_unique(this); +} + +std::shared_ptr QueryBatchKeyValueReader::GetReaderMetrics() const { + return reader_->GetReaderMetrics(); +} + +void QueryBatchKeyValueReader::Close() { + values_.reset(); + sequences_.reset(); + row_kinds_.reset(); + key_context_.reset(); + value_context_.reset(); + reader_->Close(); +} + +Result> CreateMemoryReaders( + const std::shared_ptr& split, const RealtimePartitionBucketView& memory, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& context, + const std::shared_ptr& memory_pool) { + arrow::FieldVector requested_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; + requested_fields.insert(requested_fields.end(), value_schema->fields().begin(), + value_schema->fields().end()); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(requested_fields), c_schema.get())); + ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, + memory.store->CreateQueryReaders( + memory.read_view, split->CommittedEndOffset(), query_context)); + if (batch_readers.empty()) { + return Status::Invalid("PK real-time store returned no query readers for active memory"); + } + std::vector result; + result.reserve(batch_readers.size()); + for (std::unique_ptr& reader : batch_readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null query reader"); + } + std::shared_ptr min_key; + std::shared_ptr max_key; + if (auto* provider = dynamic_cast(reader.get())) { + min_key = provider->GetMinKey(); + max_key = provider->GetMaxKey(); + } + result.push_back( + AdditionalKeyValueReader{std::make_unique( + std::move(reader), key_schema, value_schema, memory_pool), + std::move(min_key), std::move(max_key)}); + } + return result; +} + +} // namespace KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, @@ -75,6 +248,11 @@ void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { + std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); + if (realtime_split) { + return CreateRealtimeReader(realtime_split, /*release_ticket=*/true); + } + std::shared_ptr dispatch_split = split; if (auto indexed_split = std::dynamic_pointer_cast(split)) { PAIMON_RETURN_NOT_OK(indexed_split->Validate()); @@ -126,8 +304,94 @@ Result> KeyValueTableRead::CreateReader( return Status::Invalid("create reader failed, not read match with data split."); } +Result> KeyValueTableRead::CreateReader( + const std::vector>& splits) { + std::vector> readers; + readers.reserve(splits.size()); + std::vector> realtime_splits; + for (const std::shared_ptr& split : splits) { + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(split); + if (realtime_split) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + CreateRealtimeReader(realtime_split, /*release_ticket=*/false)); + readers.push_back(std::move(reader)); + realtime_splits.push_back(std::move(realtime_split)); + } else { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); + readers.push_back(std::move(reader)); + } + } + + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + if (!realtime_splits.empty()) { + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + for (const std::shared_ptr& realtime_split : realtime_splits) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + } + return result; +} + +Result> KeyValueTableRead::CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket) { + if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { + return Status::Invalid("unsupported real-time split version"); + } + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(RealtimePartitionBucketView memory, + realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); + const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), + realtime_split->Bucket()); + if (memory.partition_bucket != expected_partition_bucket) { + return Status::Invalid("real-time read-view ticket belongs to another partition-bucket"); + } + const std::optional memory_range = memory.read_view->GetOffsetRange(); + if (!memory_range || memory_range->end != realtime_split->MemoryEndOffset()) { + return Status::Invalid("real-time read-view ticket does not match the split offset range"); + } + for (const std::unique_ptr& read : split_reads_) { + auto* merge_read = dynamic_cast(read.get()); + if (merge_read) { + PAIMON_ASSIGN_OR_RAISE( + std::vector memory_readers, + CreateMemoryReaders(realtime_split, memory, merge_read->GetKeySchema(), + merge_read->GetValueSchema(), context_, GetMemoryPool())); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + merge_read->CreateRealtimeReader(realtime_split->DiskSplits(), + std::move(memory_readers))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, + RealtimeReader::Create(memory.read_view, std::move(reader))); + if (release_ticket) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + return std::unique_ptr(std::move(realtime_reader)); + } + } + return Status::Invalid("create reader failed, merge file split read not found"); +} + Result> KeyValueTableRead::CreateCountReader( const std::vector>& splits) { + for (const std::shared_ptr& split : splits) { + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented( + "CreateCountReader does not support process-local real-time splits"); + } + } if (context_->GetPredicate() != nullptr) { return Status::NotImplemented( "CreateCountReader with predicate pushdown is not supported yet"); diff --git a/src/paimon/core/table/source/key_value_table_read.h b/src/paimon/core/table/source/key_value_table_read.h index d6a1c83d3..6824ae59e 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -35,6 +35,7 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +class RealtimeSplit; class KeyValueTableRead : public TableRead { public: @@ -45,6 +46,9 @@ class KeyValueTableRead : public TableRead { Result> CreateReader(const std::shared_ptr& split) override; + Result> CreateReader( + const std::vector>& splits) override; + Result> CreateCountReader( const std::vector>& splits) override; @@ -57,6 +61,9 @@ class KeyValueTableRead : public TableRead { const std::shared_ptr& memory_pool, const std::shared_ptr& executor); + Result> CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket); + std::vector> split_reads_; std::shared_ptr path_factory_; std::shared_ptr context_; diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index c275208c5..1b496d8a1 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -108,7 +108,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl for (const std::shared_ptr& split : disk_splits) { std::shared_ptr data_split = std::dynamic_pointer_cast(split); if (!data_split) { - return Status::Invalid("real-time append scan requires process-local data splits"); + return Status::Invalid("real-time scan requires process-local data splits"); } std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h index 7d036d420..959203ca4 100644 --- a/src/paimon/core/table/source/realtime_table_scan.h +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -35,7 +35,7 @@ class FileSystem; class ScanFilter; class SnapshotManager; -/// Adds process-local memory splits to a normal append-table batch scan. +/// Adds process-local memory splits to a normal data-table batch scan. class RealtimeTableScan : public TableScan { public: RealtimeTableScan(std::unique_ptr&& disk_scan, diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 2dda955ac..b12e59a84 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -41,6 +41,7 @@ #include "paimon/core/operation/data_evolution_file_store_scan.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/primary_key_realtime_options.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/schema_validation.h" @@ -225,15 +226,15 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& if (!core_options.RealtimeEnabled()) { return Status::Invalid("real-time scan requires realtime.enabled=true"); } - if (!table_schema.PrimaryKeys().empty()) { - return Status::Invalid("real-time union read currently supports append tables only"); - } if (core_options.GetBucket() <= 0) { return Status::Invalid("real-time union read requires fixed bucket mode"); } if (core_options.DataEvolutionEnabled()) { return Status::Invalid("real-time union read does not support data evolution"); } + if (!table_schema.PrimaryKeys().empty()) { + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options)); + } if (context.IsStreamingMode()) { return Status::Invalid("real-time union read currently supports batch scans only"); } From 03ede3529c15a8f36fecf1393cbc2580eda5e28e Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:34:33 +0800 Subject: [PATCH 03/24] test(realtime): cover primary-key realtime lifecycle Cover PK write and read, recovery, external compaction, supported concurrency, writer handoff, ticket lifecycle, plugin contracts, rolling files, and multi-partition and bucket restore. --- .../operation/key_value_file_store_write.cpp | 35 +- test/inte/realtime_write_inte_test.cpp | 1005 ++++++++++++++++- 2 files changed, 977 insertions(+), 63 deletions(-) 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 4456ee1c2..e94c45a15 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -90,20 +90,6 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( } } -Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { - if (!realtime_context_) { - return Status::Invalid("refresh committed snapshot requires a real-time writer"); - } - PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); - PAIMON_ASSIGN_OR_RAISE( - RealtimeOffsetMap committed_offsets, - RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), - options_.GetFileSystem())); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(realtime_context_)); - return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets); -} - Result> KeyValueFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { PAIMON_ASSIGN_OR_RAISE( @@ -139,6 +125,7 @@ Result> KeyValueFileStoreWrite::CreateWriter( Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); std::map partition_map; int64_t materialized_max_sequence_number = restore_max_seq_number; + std::shared_ptr compact_manager; if (realtime_context_) { std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, @@ -147,15 +134,11 @@ Result> KeyValueFileStoreWrite::CreateWriter( std::map(partition_values.begin(), partition_values.end()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); - const RealtimePartitionBucket partition_bucket(partition_map, bucket); materialized_max_sequence_number = realtime_context_impl->GetMaterializedMaxSequenceNumber( - partition_bucket, restore_max_seq_number); + RealtimePartitionBucket(partition_map, bucket), restore_max_seq_number); if (materialized_max_sequence_number == std::numeric_limits::max()) { return Status::Invalid("PK sequence number has reached INT64_MAX"); } - } - std::shared_ptr compact_manager; - if (realtime_context_) { compact_manager = std::make_shared(); } else { auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); @@ -181,6 +164,20 @@ Result> KeyValueFileStoreWrite::CreateWriter( writer, options_.ToMap(), pool_, materialized_max_sequence_number); } +Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + if (!realtime_context_) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), + options_.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context_)); + return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets); +} + Status KeyValueFileStoreWrite::Close() { PAIMON_RETURN_NOT_OK(AbstractFileStoreWrite::Close()); compact_manager_factory_->Close(); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 6298137ea..f18c3f1e4 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +43,11 @@ #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/commit_context.h" +#include "paimon/common/factories/io_hook.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/realtime/realtime_context_impl.h" @@ -59,6 +64,7 @@ #include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" #include "paimon/reader/count_reader.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" #include "paimon/record_batch.h" @@ -71,6 +77,308 @@ #include "paimon/write_context.h" namespace paimon::test { +namespace { + +class BlockingState { + public: + void Block() { + std::unique_lock lock(mutex_); + entered_ = true; + entered_cv_.notify_all(); + release_cv_.wait(lock, [this]() { return released_; }); + } + + bool WaitUntilBlocked() { + std::unique_lock lock(mutex_); + return entered_cv_.wait_for(lock, std::chrono::seconds(30), [this]() { return entered_; }); + } + + void Release() { + std::lock_guard lock(mutex_); + released_ = true; + release_cv_.notify_all(); + } + + private: + std::mutex mutex_; + std::condition_variable entered_cv_; + std::condition_variable release_cv_; + bool entered_ = false; + bool released_ = false; +}; + +class BlockingBatchReader final : public BatchReader { + public: + BlockingBatchReader(std::unique_ptr&& reader, + const std::shared_ptr& state) + : reader_(std::move(reader)), state_(state) {} + + Result NextBatch() override { + if (!blocked_) { + blocked_ = true; + state_->Block(); + } + return reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr state_; + bool blocked_ = false; +}; + +class BlockingRealtimeStore final : public RealtimeStore { + public: + BlockingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr& state) + : delegate_(delegate), state_(state) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + if (!readers.empty()) { + readers[0] = std::make_unique(std::move(readers[0]), state_); + } + return readers; + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + std::shared_ptr state_; +}; + +class BlockingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit BlockingRealtimeStoreFactory(const std::shared_ptr& state) + : state_(state) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate, state_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr state_; +}; + +class TrackingRealtimeReadView final : public RealtimeReadView { + public: + explicit TrackingRealtimeReadView(std::shared_ptr delegate) + : delegate_(std::move(delegate)) {} + + std::optional GetOffsetRange() const override { + return delegate_->GetOffsetRange(); + } + + const std::shared_ptr& Delegate() const { + return delegate_; + } + + private: + std::shared_ptr delegate_; +}; + +class ReadViewCheckingBatchReader final : public BatchReader { + public: + ReadViewCheckingBatchReader(std::unique_ptr delegate, + std::weak_ptr read_view) + : delegate_(std::move(delegate)), read_view_(std::move(read_view)) {} + + Result NextBatch() override { + if (read_view_.expired()) { + return Status::Invalid("real-time read view was released before reader completion"); + } + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::weak_ptr read_view_; +}; + +class QueryTrackingRealtimeStore final : public RealtimeStore { + public: + QueryTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : delegate_(delegate), saw_query_predicate_(saw_query_predicate), query_view_(query_view) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate_view, + delegate_->AcquireReadView()); + return std::shared_ptr( + std::make_shared(delegate_view)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + if (context.predicate) { + saw_query_predicate_->store(true, std::memory_order_release); + } + *query_view_ = view; + std::shared_ptr tracking_view = + std::dynamic_pointer_cast(view); + if (!tracking_view) { + return Status::Invalid("query tracking store received an unexpected read view"); + } + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + delegate_->CreateQueryReaders(tracking_view->Delegate(), offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), view); + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +class QueryTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit QueryTrackingRealtimeStoreFactory( + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : saw_query_predicate_(saw_query_predicate), query_view_(query_view) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr(std::make_shared( + delegate, saw_query_predicate_, query_view_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +class InvalidReaderRealtimeStore final : public RealtimeStore { + public: + explicit InvalidReaderRealtimeStore(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr&) override { + std::vector> readers; + readers.push_back(nullptr); + return readers; + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { + return std::vector>(); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; +}; + +class InvalidReaderRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate)); + } + + private: + ArrowRealtimeStoreFactory delegate_; +}; + +} // namespace namespace { @@ -219,6 +527,18 @@ class RealtimeWriteInteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } + void CreatePkTable(const std::vector& partition_keys = {}) const { + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema_, c_schema.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, + Catalog::Create(dir_->Str(), options_)); + ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false)); + std::vector primary_keys = partition_keys; + primary_keys.push_back("id"); + ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), c_schema.get(), partition_keys, + primary_keys, options_, /*ignore_if_exists=*/false)); + } + Result> CreateRealtimeWriter( const std::shared_ptr& realtime_context) const { WriteContextBuilder builder(table_path_, commit_user_); @@ -240,6 +560,12 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> MakeBatch(const std::vector& rows, bool partitioned, int32_t bucket) const { + return MakeBatch(rows, partitioned, bucket, /*row_kinds=*/{}); + } + + Result> MakeBatch( + const std::vector& rows, bool partitioned, int32_t bucket, + const std::vector& row_kinds) const { if (rows.empty()) { return Status::Invalid("cannot create an empty test batch"); } @@ -263,6 +589,7 @@ class RealtimeWriteInteTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); RecordBatchBuilder builder(&c_array); + builder.SetRowKinds(row_kinds); if (partitioned) { builder.SetPartition({{"pt", partition}}); } @@ -426,6 +753,16 @@ class RealtimeWriteInteTest : public ::testing::Test { return commit->Expire(); } + Status CommitMessages(const std::vector>& messages, + int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, commit_user_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + return commit->Commit(messages, commit_identifier); + } + Result> CreatePlan( const std::shared_ptr& realtime_context, const std::shared_ptr& predicate) const { @@ -573,6 +910,75 @@ class RealtimeWriteInteTest : public ::testing::Test { return memory_usage; } + Result> ReadPkSequences( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + realtime_context_impl->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected one PK real-time read view"); + } + auto read_schema = std::make_unique(); + arrow::FieldVector requested_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; + requested_fields.insert(requested_fields.end(), schema_->fields().begin(), + schema_->fields().end()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(requested_fields), read_schema.get())); + ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + views[0].store->CreateQueryReaders(views[0].read_view, + /*offset_begin=*/0, query_context)); + std::vector sequences; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr imported, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(imported); + if (!values) { + return Status::Invalid("PK query reader did not return a StructArray"); + } + std::shared_ptr sequence_array = + std::dynamic_pointer_cast( + values->GetFieldByName(SpecialFields::SequenceNumber().Name())); + if (!sequence_array) { + return Status::Invalid("PK query reader did not return sequence numbers"); + } + for (int64_t row = 0; row < sequence_array->length(); ++row) { + sequences.push_back(sequence_array->Value(row)); + } + } + reader->Close(); + } + return sequences; + } + + static std::vector> NewFiles( + const std::vector& progresses) { + std::vector> files; + for (const RealtimeCommitProgress& progress : progresses) { + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + if (!message) { + continue; + } + const std::vector>& new_files = + message->GetNewFilesIncrement().NewFiles(); + files.insert(files.end(), new_files.begin(), new_files.end()); + } + return files; + } + static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { std::vector seen(static_cast(total_rows), false); int64_t max_id = -1; @@ -627,7 +1033,6 @@ class RealtimeWriteInteTest : public ::testing::Test { options_[Options::PARTITION_GENERATE_LEGACY_NAME] = legacy_partition_name_enabled ? "true" : "false"; CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -666,6 +1071,57 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_OK(writer->Close()); } + void CheckVectorReaderRetry(bool primary_key) { + if (primary_key) { + CreatePkTable(/*partition_keys=*/{"pt"}); + } else { + CreateTable(/*partition_keys=*/{"pt"}); + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, + MakeBatch(p0_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p0_batch))); + std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, + MakeBatch(p1_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p1_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(2, plan->Splits().size()); + + std::vector> invalid_splits = plan->Splits(); + std::shared_ptr second_split = + std::dynamic_pointer_cast(invalid_splits[1]); + ASSERT_NE(nullptr, second_split); + std::vector> second_disk_splits = second_split->DiskSplits(); + invalid_splits[1] = std::make_shared( + RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), + second_split->Partition(), second_split->Bucket(), std::move(second_disk_splits), + second_split->CommittedEndOffset(), second_split->MemoryEndOffset(), + second_split->OpaqueTicket()); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), + "unsupported real-time split version"); + + std::vector expected_rows = p0_rows; + expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); + } + std::unique_ptr dir_; std::string table_path_; std::string commit_user_ = "realtime_commit_user"; @@ -723,6 +1179,459 @@ TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } +TEST_F(RealtimeWriteInteTest, TestPkRead) { + CreatePkTable(); + auto saw_query_predicate = std::make_shared>(false); + auto query_view = std::make_shared>(); + auto factory = + std::make_shared(saw_query_predicate, query_view); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + MakeBatch({Row{1, "new", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(update_batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "new", "p0"}, {2, "two", "p0"}}), memory_rows); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + + std::vector second_rows = {{1, "latest", "p0"}, {2, "gone", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector union_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "latest", "p0"}, {3, "three", "p0"}}), union_rows); + + const std::string expected_payload = "new"; + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, expected_payload.data(), expected_payload.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_plan, + CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult filtered_result, + ReadPlan(filtered_plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/true)); + ASSERT_EQ(nullptr, filtered_result.data); + ASSERT_FALSE(saw_query_predicate->load(std::memory_order_acquire)); + filtered_result.reader->Close(); + filtered_result.reader.reset(); + ASSERT_OK(writer->Close()); + writer.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr lifetime_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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 reader, + table_read->CreateReader(lifetime_plan->Splits())); + ASSERT_FALSE(query_view->expired()); + + std::weak_ptr weak_context = realtime_context; + table_read.reset(); + lifetime_plan.reset(); + realtime_context.reset(); + ASSERT_TRUE(weak_context.expired()); + ASSERT_FALSE(query_view->expired()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch read_batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(read_batch)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_array, + ReadResultCollector::GetArray(std::move(read_batch))); + ASSERT_NE(nullptr, read_array); + read_array.reset(); + reader->Close(); + reader.reset(); + ASSERT_TRUE(query_view->expired()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRecovery) { + CreatePkTable(); + + WriteContextBuilder seed_builder(table_path_, commit_user_); + seed_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_context, seed_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_writer, + FileStoreWrite::Create(std::move(seed_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_batch, + MakeBatch({Row{99, "seed", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(seed_writer->Write(std::move(seed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> seed_messages, + seed_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/0)); + CommitContextBuilder seed_commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit_context, + seed_commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit, + FileStoreCommit::Create(std::move(seed_commit_context))); + ASSERT_OK(seed_commit->Commit(seed_messages)); + ASSERT_OK(seed_writer->Close()); + const std::vector mutations = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector mutation_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result> failed_prepare = + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1); + io_hook->Clear(); + ASSERT_TRUE(failed_prepare.status().IsIOError()) << failed_prepare.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failed_prepare, ReadRows()); + ASSERT_EQ((std::vector{{99, "seed", "p0"}}), rows_after_failed_prepare); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(first_writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_sequences, ReadPkSequences(first_context)); + ASSERT_EQ((std::vector{2, 3, 4}), memory_sequences); + ASSERT_OK_AND_ASSIGN(std::vector progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_EQ(memory_sequences.front(), NewFiles(progress)[0]->min_sequence_number); + ASSERT_EQ(memory_sequences.back(), NewFiles(progress)[0]->max_sequence_number); + ASSERT_OK(Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK(first_writer->Close()); + first_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_replay, ReadRows()); + ASSERT_EQ((std::vector{{1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}), + rows_after_replay); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr restart_batch, + MakeBatch({Row{4, "four", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(restart_batch))); + ASSERT_OK_AND_ASSIGN(std::vector restart_sequences, ReadPkSequences(second_context)); + ASSERT_EQ((std::vector{5}), restart_sequences); + ASSERT_OK_AND_ASSIGN(std::vector restart_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, restart_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), restart_progress[0].offset_range); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->min_sequence_number); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->max_sequence_number); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompaction) { + options_[Options::NUM_SORTED_RUNS_COMPACTION_TRIGGER] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + int64_t latest_snapshot_id = -1; + constexpr int64_t kCommitRoundsBeforeCompaction = 4; + std::set committed_file_names; + for (int64_t round = 0; round < kCommitRoundsBeforeCompaction; ++round) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{round, "value-" + std::to_string(round), "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(round)); + ASSERT_EQ(1, progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(progress[0].commit_message); + ASSERT_NE(nullptr, message); + ASSERT_TRUE(message->GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, NewFiles(progress).size()); + committed_file_names.insert(NewFiles(progress)[0]->file_name); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(progress, round)); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr next_batch, + MakeBatch({Row{4, "value-4", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(next_batch))); + + WriteContextBuilder compact_builder(table_path_, commit_user_); + compact_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_context, compact_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_writer, + FileStoreWrite::Create(std::move(compact_context))); + ASSERT_OK(compact_writer->Compact(/*partition=*/{}, /*bucket=*/0, + /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> compact_messages, + compact_writer->PrepareCommit(/*wait_compaction=*/true, /*commit_identifier=*/4)); + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact_message = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact_message); + ASSERT_TRUE(compact_message->GetNewFilesIncrement().IsEmpty()); + ASSERT_EQ(kCommitRoundsBeforeCompaction, + compact_message->GetCompactIncrement().CompactBefore().size()); + std::set compacted_file_names; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactBefore()) { + compacted_file_names.insert(file->file_name); + } + ASSERT_EQ(committed_file_names, compacted_file_names); + ASSERT_FALSE(compact_message->GetCompactIncrement().CompactAfter().empty()); + ASSERT_OK(CommitMessages(compact_messages, /*commit_identifier=*/4)); + ASSERT_OK(compact_writer->Close()); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional compact_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(compact_snapshot); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot->GetCommitKind()); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(4, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK(writer->RefreshCommittedSnapshot(compact_snapshot->Id())); + ASSERT_OK_AND_ASSIGN(std::vector compacted_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, + {1, "value-1", "p0"}, + {2, "value-2", "p0"}, + {3, "value-3", "p0"}, + {4, "value-4", "p0"}}), + compacted_rows); + + constexpr int64_t kCommitRoundsAfterCompaction = 2; + for (int64_t round = 0; round < kCommitRoundsAfterCompaction; ++round) { + if (round > 0) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeBatch({Row{4 + round, "value-" + std::to_string(4 + round), "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + } + const int64_t commit_identifier = 5 + round; + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(4 + round, 5 + round), progress[0].offset_range); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + } + ASSERT_OK_AND_ASSIGN(offsets, ReadCommittedOffsets()); + ASSERT_EQ(6, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, + {1, "value-1", "p0"}, + {2, "value-2", "p0"}, + {3, "value-3", "p0"}, + {4, "value-4", "p0"}, + {5, "value-5", "p0"}}), + final_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + constexpr int64_t kRowCount = 20; + constexpr int32_t kReaderCount = 2; + std::atomic writer_done{false}; + std::atomic control_done{false}; + std::atomic commit_count{0}; + ConcurrentTestState state; + std::vector read_counts(kReaderCount, 0); + + std::thread write_thread([&]() { + state.WaitForStart(); + for (int64_t id = 0; id < kRowCount && !state.ShouldStop(); ++id) { + Result> batch = + MakeBatch(MakeRows(id, /*count=*/1, /*partition=*/"p0"), + /*partitioned=*/false); + if (state.RecordErrorIfNotOk(batch) || + state.RecordErrorIfNotOk(writer->Write(std::move(batch).value()))) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + writer_done.store(true, std::memory_order_release); + }); + + std::thread control_thread([&]() { + state.WaitForStart(); + int64_t commit_identifier = 0; + do { + Result> progress = + writer->PrepareCommitWithProgress(commit_identifier); + if (state.RecordErrorIfNotOk(progress)) { + break; + } + if (!progress.value().empty()) { + Result snapshot = Commit(progress.value(), commit_identifier++); + if (state.RecordErrorIfNotOk(snapshot) || + state.RecordErrorIfNotOk(writer->RefreshCommittedSnapshot(snapshot.value()))) { + break; + } + commit_count.fetch_add(1, std::memory_order_relaxed); + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } while (!writer_done.load(std::memory_order_acquire) && !state.ShouldStop()); + if (!state.ShouldStop()) { + Result> progress = + writer->PrepareCommitWithProgress(commit_identifier); + if (!state.RecordErrorIfNotOk(progress) && !progress.value().empty()) { + Result snapshot = Commit(progress.value(), commit_identifier); + if (!state.RecordErrorIfNotOk(snapshot) && + !state.RecordErrorIfNotOk(writer->RefreshCommittedSnapshot(snapshot.value()))) { + commit_count.fetch_add(1, std::memory_order_relaxed); + } + } + } + control_done.store(true, std::memory_order_release); + }); + + std::vector read_threads; + read_threads.reserve(kReaderCount); + for (int32_t reader_index = 0; reader_index < kReaderCount; ++reader_index) { + read_threads.emplace_back([&, reader_index]() { + state.WaitForStart(); + while (!control_done.load(std::memory_order_acquire) && !state.ShouldStop()) { + Result> rows = ReadRows(realtime_context); + ++read_counts[reader_index]; + if (state.RecordErrorIfNotOk(rows) || + state.RecordErrorIfNotOk(ValidateReadPrefix(rows.value(), kRowCount))) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + }); + } + + state.StartWhenReady(/*worker_count=*/2 + kReaderCount); + write_thread.join(); + control_thread.join(); + for (std::thread& read_thread : read_threads) { + read_thread.join(); + } + + ASSERT_TRUE(state.Errors().empty()) << (state.Errors().empty() ? "" : state.Errors().front()); + ASSERT_GT(commit_count.load(), 0); + for (int32_t read_count : read_counts) { + ASSERT_GT(read_count, 0); + } + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); + ASSERT_EQ(kRowCount, static_cast(final_rows.size())); + ASSERT_OK(ValidateReadPrefix(final_rows, kRowCount)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkWriteDuringPrepare) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + + Result> prepare_result = + Status::Invalid("prepare did not run"); + std::thread prepare_thread( + [&]() { prepare_result = writer->PrepareCommitWithProgress(/*commit_identifier=*/0); }); + const bool prepare_blocked = state->WaitUntilBlocked(); + if (!prepare_blocked) { + state->Release(); + prepare_thread.join(); + ASSERT_TRUE(prepare_blocked); + } + + std::promise write_promise; + std::future write_future = write_promise.get_future(); + std::thread write_thread([&]() { + Result> batch = + MakeBatch({Row{2, "two", "p0"}}, /*partitioned=*/false); + if (!batch.ok()) { + write_promise.set_value(batch.status()); + return; + } + write_promise.set_value(writer->Write(std::move(batch).value())); + }); + const bool write_completed = + write_future.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + state->Release(); + prepare_thread.join(); + write_thread.join(); + + ASSERT_TRUE(write_completed); + ASSERT_OK(write_future.get()); + ASSERT_OK(prepare_result); + ASSERT_EQ(1, prepare_result.value().size()); + ASSERT_EQ(OffsetRange(0, 1), prepare_result.value()[0].offset_range); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_progress.size()); + ASSERT_EQ(OffsetRange(1, 2), second_progress[0].offset_range); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { + CreatePkTable(); + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + ASSERT_NOK_WITH_MSG(ReadRows(realtime_context), + "PK real-time store returned no query readers for active memory"); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PK real-time store returned a null commit reader"); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { options_[Options::TARGET_FILE_ROW_NUM] = "10"; CreateTable(/*partition_keys=*/{}); @@ -1235,50 +2144,12 @@ TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTick ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestVectorReaderFailurePreservesEarlierSplitTicket) { - CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, - MakeBatch(p0_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p0_batch))); - std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, - MakeBatch(p1_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p1_batch))); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, - CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_EQ(2, plan->Splits().size()); - - std::vector> invalid_splits = plan->Splits(); - std::shared_ptr second_split = - std::dynamic_pointer_cast(invalid_splits[1]); - ASSERT_NE(nullptr, second_split); - std::vector> second_disk_splits = second_split->DiskSplits(); - invalid_splits[1] = std::make_shared( - RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), second_split->Partition(), - second_split->Bucket(), std::move(second_disk_splits), second_split->CommittedEndOffset(), - second_split->MemoryEndOffset(), second_split->OpaqueTicket()); - - ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(options_) - .SetReadFieldNames({"id", "payload", "pt"}) - .WithRealtimeContext(realtime_context) - .WithMemoryPool(pool_); - 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_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), - "unsupported real-time split version"); +TEST_F(RealtimeWriteInteTest, TestAppendVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/false); +} - std::vector expected_rows = p0_rows; - expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); - ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); +TEST_F(RealtimeWriteInteTest, TestPkVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/true); } TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { @@ -1351,6 +2222,52 @@ TEST_F(RealtimeWriteInteTest, TestCloseWriterAllowsContextReuseByLaterWriter) { ASSERT_OK(second_writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(realtime_context)); + std::vector first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(OffsetRange(0, 3), commits[0].offset_range); + ASSERT_EQ(1, NewFiles(commits).size()); + ASSERT_EQ(0, NewFiles(commits)[0]->min_sequence_number); + ASSERT_EQ(2, NewFiles(commits)[0]->max_sequence_number); + ASSERT_OK(first_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(realtime_context)); + std::vector second_rows = { + Row{0, "updated-0", "p0"}, + Row{3, "value-3", "p0"}, + }; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_commits.size()); + ASSERT_EQ(OffsetRange(3, 5), second_commits[0].offset_range); + ASSERT_EQ(1, NewFiles(second_commits).size()); + ASSERT_EQ(3, NewFiles(second_commits)[0]->min_sequence_number); + ASSERT_EQ(4, NewFiles(second_commits)[0]->max_sequence_number); + + commits.push_back(std::move(second_commits[0])); + ASSERT_OK(Commit(commits, /*commit_identifier=*/1)); + std::vector expected_rows = first_rows; + expected_rows[0] = second_rows[0]; + expected_rows.push_back(second_rows[1]); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(second_writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, From 3e256308f58cd9fd3b04147c17cf5f9e83f748ea Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:01:40 +0800 Subject: [PATCH 04/24] refactor(realtime): consolidate PK state and validation --- src/paimon/CMakeLists.txt | 2 - .../core/operation/file_store_write.cpp | 3 +- .../operation/key_value_file_store_write.cpp | 30 ++++++---- .../realtime/primary_key_realtime_options.cpp | 58 ------------------- .../realtime/primary_key_realtime_options.h | 31 ---------- .../primary_key_realtime_options_test.cpp | 56 ------------------ .../core/realtime/realtime_context_impl.cpp | 27 ++++----- .../core/realtime/realtime_context_impl.h | 4 +- .../core/realtime/realtime_context_test.cpp | 38 ++++++++++++ .../realtime/realtime_primary_key_writer.cpp | 41 ++----------- .../realtime/realtime_primary_key_writer.h | 13 ++--- src/paimon/core/table/source/table_scan.cpp | 4 +- .../core/utils/primary_key_table_utils.cpp | 32 ++++++++++ .../core/utils/primary_key_table_utils.h | 3 + .../utils/primary_key_table_utils_test.cpp | 26 +++++++++ 15 files changed, 144 insertions(+), 224 deletions(-) delete mode 100644 src/paimon/core/realtime/primary_key_realtime_options.cpp delete mode 100644 src/paimon/core/realtime/primary_key_realtime_options.h delete mode 100644 src/paimon/core/realtime/primary_key_realtime_options_test.cpp diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 69deab92f..b2e1e6617 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -379,7 +379,6 @@ set(PAIMON_CORE_SRCS core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp core/realtime/primary_key_realtime_store.cpp - core/realtime/primary_key_realtime_options.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp @@ -784,7 +783,6 @@ if(PAIMON_BUILD_TESTS) core/memory/writer_memory_manager_test.cpp core/realtime/arrow_realtime_store_test.cpp core/realtime/primary_key_realtime_store_test.cpp - core/realtime/primary_key_realtime_options_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index fb83c254c..f216476bd 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -36,7 +36,6 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" -#include "paimon/core/realtime/primary_key_realtime_options.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" @@ -198,7 +197,7 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options)); + PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); if (ignore_previous_files) { return Status::NotImplemented( "PK realtime v1 requires restore from the latest snapshot"); 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 e94c45a15..492161cf8 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -124,19 +124,28 @@ Result> KeyValueFileStoreWrite::CreateWriter( std::shared_ptr levels, Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); std::map partition_map; - int64_t materialized_max_sequence_number = restore_max_seq_number; + int64_t initial_max_sequence_number = restore_max_seq_number; std::shared_ptr compact_manager; + std::shared_ptr realtime_context_impl; + std::optional realtime_store_state; if (realtime_context_) { std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, file_store_path_factory_->GeneratePartitionVector(partition)); partition_map = std::map(partition_values.begin(), partition_values.end()); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(realtime_context_)); - materialized_max_sequence_number = realtime_context_impl->GetMaterializedMaxSequenceNumber( - RealtimePartitionBucket(partition_map, bucket), restore_max_seq_number); - if (materialized_max_sequence_number == std::numeric_limits::max()) { + PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, c_write_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ + std::move(c_write_schema), options_.ToMap(), pool_, partition_map, bucket, + PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys, + restore_max_seq_number}})); + realtime_store_state = std::move(store_state); + initial_max_sequence_number = realtime_store_state->initial_max_sequence_number.value(); + if (initial_max_sequence_number == std::numeric_limits::max()) { return Status::Invalid("PK sequence number has reached INT64_MAX"); } compact_manager = std::make_shared(); @@ -150,18 +159,15 @@ Result> KeyValueFileStoreWrite::CreateWriter( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( - materialized_max_sequence_number, trimmed_primary_keys, data_file_path_factory, + initial_max_sequence_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, options_, compact_manager, realtime_context_ ? nullptr : io_manager_, enable_multi_thread_spill_, pool_)); if (!realtime_context_) { return std::shared_ptr(std::move(writer)); } - auto c_write_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, c_write_schema.get())); - return RealtimePrimaryKeyWriter::Create( - partition_map, bucket, std::move(c_write_schema), trimmed_primary_keys, realtime_context_, - writer, options_.ToMap(), pool_, materialized_max_sequence_number); + return RealtimePrimaryKeyWriter::Create(partition_map, bucket, schema_, realtime_context_impl, + writer, pool_, realtime_store_state.value()); } Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { diff --git a/src/paimon/core/realtime/primary_key_realtime_options.cpp b/src/paimon/core/realtime/primary_key_realtime_options.cpp deleted file mode 100644 index e9779a59e..000000000 --- a/src/paimon/core/realtime/primary_key_realtime_options.cpp +++ /dev/null @@ -1,58 +0,0 @@ -/* - * 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/realtime/primary_key_realtime_options.h" - -#include "paimon/core/core_options.h" - -namespace paimon { - -Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { - if (options.GetBucket() <= 0) { - return Status::NotImplemented("PK realtime v1 requires fixed buckets"); - } - if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { - return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); - } - if (options.DataEvolutionEnabled()) { - return Status::NotImplemented("PK realtime v1 does not support data evolution"); - } - if (!options.GetFieldsSequenceGroups().empty()) { - return Status::NotImplemented("PK realtime v1 does not support sequence groups"); - } - if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || - options.AggregationRemoveRecordOnDelete() || - !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { - return Status::NotImplemented("PK realtime v1 requires default delete behavior"); - } - if (!options.GetSequenceField().empty()) { - return Status::NotImplemented("PK realtime v1 does not support sequence.field"); - } - if (!options.SequenceFieldSortOrderIsAscending()) { - return Status::NotImplemented( - "PK realtime v1 supports only ascending sequence.field.sort-order"); - } - if (options.NeedLookup() || options.DeletionVectorsEnabled() || - options.GetChangelogProducer() != ChangelogProducer::NONE) { - return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); - } - return Status::OK(); -} - -} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_options.h b/src/paimon/core/realtime/primary_key_realtime_options.h deleted file mode 100644 index a16d35778..000000000 --- a/src/paimon/core/realtime/primary_key_realtime_options.h +++ /dev/null @@ -1,31 +0,0 @@ -/* - * 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 "paimon/status.h" - -namespace paimon { - -class CoreOptions; - -/// Validates the table options supported by the in-memory PK realtime V1 path. -Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options); - -} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_options_test.cpp b/src/paimon/core/realtime/primary_key_realtime_options_test.cpp deleted file mode 100644 index 5d3ea7f67..000000000 --- a/src/paimon/core/realtime/primary_key_realtime_options_test.cpp +++ /dev/null @@ -1,56 +0,0 @@ -/* - * 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/realtime/primary_key_realtime_options.h" - -#include -#include -#include - -#include "paimon/core/core_options.h" -#include "paimon/defs.h" -#include "paimon/testing/utils/testharness.h" - -namespace paimon::test { - -TEST(PrimaryKeyRealtimeOptionsTest, TestSupportedOptions) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); - ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options)); -} - -TEST(PrimaryKeyRealtimeOptionsTest, TestUnsupportedOptions) { - const std::string sequence_group = - std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; - const std::vector> unsupported_options = { - {{Options::BUCKET, "0"}}, - {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, - {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, - {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, - {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, - {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, - {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, - {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, - }; - for (const std::map& option_map : unsupported_options) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); - ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options)); - } -} - -} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 0a367b2cd..6624059a6 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -82,6 +82,18 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); const RealtimePartitionBucket key(request.partition, request.bucket); + std::optional initial_max_sequence_number; + PrimaryKeyRealtimeStoreCreateConfig* primary_key_config = + std::get_if(&request.mode_config); + if (primary_key_config) { + auto [sequence_iter, inserted] = materialized_max_sequence_numbers_.emplace( + key, primary_key_config->restore_max_sequence_number); + if (!inserted && primary_key_config->restore_max_sequence_number > sequence_iter->second) { + sequence_iter->second = primary_key_config->restore_max_sequence_number; + } + initial_max_sequence_number = sequence_iter->second; + primary_key_config->restore_max_sequence_number = sequence_iter->second; + } int64_t initial_offset = 0; auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { @@ -114,7 +126,7 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset}; + return RealtimeStoreState{iter->second, initial_offset, initial_max_sequence_number}; } Result> store_result = factory_->Create(std::move(request)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); @@ -122,18 +134,7 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); } - return RealtimeStoreState{std::move(store), initial_offset}; -} - -int64_t RealtimeContextImpl::GetMaterializedMaxSequenceNumber( - const RealtimePartitionBucket& partition_bucket, int64_t restored_max_sequence_number) { - std::lock_guard lock(mutex_); - auto [iter, inserted] = - materialized_max_sequence_numbers_.emplace(partition_bucket, restored_max_sequence_number); - if (!inserted && restored_max_sequence_number > iter->second) { - iter->second = restored_max_sequence_number; - } - return iter->second; + return RealtimeStoreState{std::move(store), initial_offset, initial_max_sequence_number}; } void RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 45d07deeb..f4cd3866e 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -47,6 +47,7 @@ class MemoryPool; struct RealtimeStoreState { std::shared_ptr store; int64_t initial_offset; + std::optional initial_max_sequence_number; }; struct RealtimePartitionBucketView { @@ -67,9 +68,6 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); - int64_t GetMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, - int64_t restored_max_sequence_number); - void AdvanceMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number); diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 33701afac..b4d2c6718 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -129,6 +129,15 @@ Result GetOrCreateAppendStore( AppendRealtimeStoreCreateConfig{StatisticsMode::NONE}}); } +Result GetOrCreatePrimaryKeyStore( + const std::shared_ptr& context, + const std::map& partition, int32_t bucket, + int64_t restore_max_sequence_number, const std::shared_ptr& memory_pool) { + return context->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ + MakeWriteSchema(), /*options=*/{}, memory_pool, partition, bucket, + PrimaryKeyRealtimeStoreCreateConfig{{"id"}, restore_max_sequence_number}}); +} + TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); @@ -138,6 +147,7 @@ TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {{"k", "v"}}, pool)); ASSERT_EQ(0, first_state.initial_offset); + ASSERT_FALSE(first_state.initial_max_sequence_number); ASSERT_OK_AND_ASSIGN( RealtimeStoreState first_again_state, GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, pool)); @@ -168,6 +178,34 @@ TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { ASSERT_EQ(1, factory->stores[2]->acquire_count); } +TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState first_state, + GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, + /*restore_max_sequence_number=*/4, GetDefaultPool())); + ASSERT_EQ(4, first_state.initial_max_sequence_number); + + const RealtimePartitionBucket partition_bucket(partition, /*bucket=*/0); + context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, /*max_sequence_number=*/8); + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState retained_state, + GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, + /*restore_max_sequence_number=*/6, GetDefaultPool())); + ASSERT_EQ(first_state.store, retained_state.store); + ASSERT_EQ(8, retained_state.initial_max_sequence_number); + + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState restored_state, + GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, + /*restore_max_sequence_number=*/10, GetDefaultPool())); + ASSERT_EQ(first_state.store, restored_state.store); + ASSERT_EQ(10, restored_state.initial_max_sequence_number); +} + TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 2ebcede82..e33f48bba 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -44,44 +44,13 @@ namespace paimon { Result> RealtimePrimaryKeyWriter::Create( const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, - const std::vector& trimmed_primary_keys, - const std::shared_ptr& realtime_context, + const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_context, const std::shared_ptr& merge_tree_writer, - const std::map& options, - const std::shared_ptr& memory_pool, int64_t restore_max_sequence_number) { - ScopeGuard schema_guard([schema = write_schema.get()]() { - if (schema && schema->release) { - ArrowSchemaRelease(schema); - } - }); - if (!realtime_context) { - return Status::Invalid("PK real-time context is null"); - } - if (!merge_tree_writer) { - return Status::Invalid("PK real-time merge-tree writer is null"); - } - if (!write_schema || !write_schema->release) { - return Status::Invalid("PK real-time write schema is null"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, - arrow::ImportSchema(write_schema.get())); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*imported_schema, write_schema.get())); - RealtimeStoreCreateRequest request{ - std::move(write_schema), - options, - memory_pool, - partition, - bucket, - PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys, restore_max_sequence_number}}; - schema_guard.Release(); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(realtime_context)); - PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, - realtime_context_impl->GetOrCreateRealtimeStore(std::move(request))); + const std::shared_ptr& memory_pool, const RealtimeStoreState& store_state) { return std::shared_ptr( - new RealtimePrimaryKeyWriter(store_state.store, merge_tree_writer, realtime_context_impl, - RealtimePartitionBucket(partition, bucket), imported_schema, + new RealtimePrimaryKeyWriter(store_state.store, merge_tree_writer, realtime_context, + RealtimePartitionBucket(partition, bucket), write_schema, store_state.initial_offset, memory_pool)); } diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index fa057e079..c1e893c85 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -24,14 +24,11 @@ #include #include #include -#include #include "paimon/core/utils/batch_writer.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" -struct ArrowSchema; - namespace arrow { class Schema; } // namespace arrow @@ -40,20 +37,18 @@ namespace paimon { class MemoryPool; class MergeTreeWriter; -class RealtimeContext; class RealtimeContextImpl; +struct RealtimeStoreState; /// Primary-key real-time writer backed by an in-memory mutation indexer. class RealtimePrimaryKeyWriter final : public BatchWriter { public: static Result> Create( const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, - const std::vector& trimmed_primary_keys, - const std::shared_ptr& realtime_context, + const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_context, const std::shared_ptr& merge_tree_writer, - const std::map& options, - const std::shared_ptr& memory_pool, int64_t restore_max_sequence_number); + const std::shared_ptr& memory_pool, const RealtimeStoreState& store_state); Status Write(std::unique_ptr&& batch) override; Result PrepareCommit(bool wait_compaction) override; diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index b12e59a84..92155de3b 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -41,7 +41,6 @@ #include "paimon/core/operation/data_evolution_file_store_scan.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" -#include "paimon/core/realtime/primary_key_realtime_options.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/schema_validation.h" @@ -64,6 +63,7 @@ #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/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/format/file_format.h" #include "paimon/realtime/realtime_context.h" @@ -233,7 +233,7 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& return Status::Invalid("real-time union read does not support data evolution"); } if (!table_schema.PrimaryKeys().empty()) { - PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options)); + PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(core_options)); } if (context.IsStreamingMode()) { return Status::Invalid("real-time union read currently supports batch scans only"); diff --git a/src/paimon/core/utils/primary_key_table_utils.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index cf72da4ae..823d48c41 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -96,4 +96,36 @@ Result> PrimaryKeyTableUtils::CreateSequenceFi options.SequenceFieldSortOrderIsAscending()); } +Status PrimaryKeyTableUtils::ValidateRealtimeOptions(const CoreOptions& options) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime v1 requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime v1 does not support data evolution"); + } + if (!options.GetFieldsSequenceGroups().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence groups"); + } + if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || + options.AggregationRemoveRecordOnDelete() || + !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { + return Status::NotImplemented("PK realtime v1 requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime v1 supports only ascending sequence.field.sort-order"); + } + if (options.NeedLookup() || options.DeletionVectorsEnabled() || + options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); + } + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index 82a108ab7..7877ee4ab 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -24,6 +24,7 @@ #include "arrow/type.h" #include "paimon/result.h" +#include "paimon/status.h" namespace arrow { class Schema; @@ -57,6 +58,8 @@ class PrimaryKeyTableUtils { static Result> CreateSequenceFieldsComparator( const std::vector& value_fields, const CoreOptions& options); + + static Status ValidateRealtimeOptions(const CoreOptions& options); }; } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 12713ca5b..072965cff 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -22,7 +22,9 @@ #include #include #include +#include #include +#include #include "arrow/type.h" #include "gtest/gtest.h" @@ -112,4 +114,28 @@ TEST(PrimaryKeyTableUtilsTest, TestCreateFirstRowMergeFunctionWithIgnoreDelete) "First row merge engine can not accept DELETE/UPDATE_BEFORE records"); } +TEST(PrimaryKeyTableUtilsTest, TestSupportedRealtimeOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); +} + +TEST(PrimaryKeyTableUtilsTest, TestUnsupportedRealtimeOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); + } +} + } // namespace paimon::test From 9f6c99d1a1098ca2c3c2898103919e9f411c1313 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 14:55:20 +0800 Subject: [PATCH 05/24] fix(read): close PK realtime query readers --- .../table/source/key_value_table_read.cpp | 9 +- test/inte/realtime_write_inte_test.cpp | 138 ++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 770caf1ca..59041b78e 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -65,6 +65,10 @@ class QueryBatchKeyValueReader final : public KeyValueRecordReader { value_schema_(value_schema), pool_(pool) {} + ~QueryBatchKeyValueReader() override { + Close(); + } + Result> NextBatch() override; std::shared_ptr GetReaderMetrics() const override; void Close() override; @@ -161,7 +165,10 @@ void QueryBatchKeyValueReader::Close() { row_kinds_.reset(); key_context_.reset(); value_context_.reset(); - reader_->Close(); + if (reader_) { + reader_->Close(); + reader_.reset(); + } } Result> CreateMemoryReaders( diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index f18c3f1e4..cee96301f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -324,6 +324,101 @@ class QueryTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { std::shared_ptr> query_view_; }; +class CloseTrackingBatchReader final : public BatchReader { + public: + CloseTrackingBatchReader(std::unique_ptr delegate, + const std::shared_ptr>& close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + close_count_->fetch_add(1, std::memory_order_release); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr> close_count_; +}; + +class CloseTrackingRealtimeStore final : public RealtimeStore { + public: + CloseTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr>& close_count, + const std::shared_ptr>& append_null_reader) + : delegate_(delegate), close_count_(close_count), append_null_reader_(append_null_reader) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateQueryReaders(view, offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), close_count_); + } + if (append_null_reader_->load(std::memory_order_acquire)) { + readers.push_back(nullptr); + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + std::shared_ptr> close_count_; + std::shared_ptr> append_null_reader_; +}; + +class CloseTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + CloseTrackingRealtimeStoreFactory(const std::shared_ptr>& close_count, + const std::shared_ptr>& append_null_reader) + : close_count_(close_count), append_null_reader_(append_null_reader) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr(std::make_shared( + delegate, close_count_, append_null_reader_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr> close_count_; + std::shared_ptr> append_null_reader_; +}; + class InvalidReaderRealtimeStore final : public RealtimeStore { public: explicit InvalidReaderRealtimeStore(const std::shared_ptr& delegate) @@ -1632,6 +1727,49 @@ TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkPluginQueryReaderCloseLifecycle) { + CreatePkTable(); + auto close_count = std::make_shared>(0); + auto append_null_reader = std::make_shared>(false); + auto factory = + std::make_shared(close_count, append_null_reader); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + auto create_reader = [&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + return table_read->CreateReader(plan->Splits()); + }; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr explicitly_closed_reader, create_reader()); + explicitly_closed_reader->Close(); + explicitly_closed_reader.reset(); + ASSERT_EQ(1, close_count->load(std::memory_order_acquire)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr destroyed_reader, create_reader()); + destroyed_reader.reset(); + ASSERT_EQ(2, close_count->load(std::memory_order_acquire)); + + append_null_reader->store(true, std::memory_order_release); + ASSERT_NOK_WITH_MSG(create_reader(), "PK real-time store returned a null query reader"); + ASSERT_EQ(3, close_count->load(std::memory_order_acquire)); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { options_[Options::TARGET_FILE_ROW_NUM] = "10"; CreateTable(/*partition_keys=*/{}); From 15d7c913d6055006698980102f7df8d8a48c6e8e Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:41:07 +0800 Subject: [PATCH 06/24] fix(realtime): close rejected plugin readers --- .../realtime/realtime_primary_key_writer.cpp | 7 + .../table/source/key_value_table_read.cpp | 7 + test/inte/realtime_write_inte_test.cpp | 132 ++++++++++++++---- 3 files changed, 119 insertions(+), 27 deletions(-) diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index e33f48bba..65bcebcad 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -117,6 +117,13 @@ Status RealtimePrimaryKeyWriter::FlushSegment( const std::shared_ptr& segment) { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, realtime_store_->CreateCommitReaders(segment)); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::unique_ptr& reader : readers) { if (!reader) { return Status::Invalid("PK real-time store returned a null commit reader"); diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 59041b78e..9b5f6ee83 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -190,6 +190,13 @@ Result> CreateMemoryReaders( PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, memory.store->CreateQueryReaders( memory.read_view, split->CommittedEndOffset(), query_context)); + ScopeGuard reader_guard([&batch_readers]() { + for (const std::unique_ptr& reader : batch_readers) { + if (reader) { + reader->Close(); + } + } + }); if (batch_readers.empty()) { return Status::Invalid("PK real-time store returned no query readers for active memory"); } diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index cee96301f..e6000561f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -348,12 +348,20 @@ class CloseTrackingBatchReader final : public BatchReader { std::shared_ptr> close_count_; }; +struct CloseTrackingReaderState { + std::shared_ptr> query_close_count = + std::make_shared>(0); + std::shared_ptr> commit_close_count = + std::make_shared>(0); + int32_t query_null_index = -1; + int32_t commit_null_index = -1; +}; + class CloseTrackingRealtimeStore final : public RealtimeStore { public: CloseTrackingRealtimeStore(const std::shared_ptr& delegate, - const std::shared_ptr>& close_count, - const std::shared_ptr>& append_null_reader) - : delegate_(delegate), close_count_(close_count), append_null_reader_(append_null_reader) {} + const std::shared_ptr& state) + : delegate_(delegate), state_(state) {} Status Write(RealtimeWriteBatch&& batch) override { return delegate_->Write(std::move(batch)); @@ -365,7 +373,14 @@ class CloseTrackingRealtimeStore final : public RealtimeStore { Result>> CreateCommitReaders( const std::shared_ptr& segment) override { - return delegate_->CreateCommitReaders(segment); + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), + state_->commit_close_count); + } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->commit_null_index, &readers)); + return readers; } Result> AcquireReadView() override { @@ -378,11 +393,10 @@ class CloseTrackingRealtimeStore final : public RealtimeStore { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, delegate_->CreateQueryReaders(view, offset_begin, context)); for (std::unique_ptr& reader : readers) { - reader = std::make_unique(std::move(reader), close_count_); - } - if (append_null_reader_->load(std::memory_order_acquire)) { - readers.push_back(nullptr); + reader = std::make_unique(std::move(reader), + state_->query_close_count); } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->query_null_index, &readers)); return readers; } @@ -395,28 +409,38 @@ class CloseTrackingRealtimeStore final : public RealtimeStore { } private: + static Status InsertNullReader(int32_t index, + std::vector>* readers) { + if (index < 0) { + return Status::OK(); + } + if (index > static_cast(readers->size())) { + return Status::Invalid("null reader index exceeds reader count"); + } + readers->insert(readers->begin() + index, nullptr); + return Status::OK(); + } + std::shared_ptr delegate_; - std::shared_ptr> close_count_; - std::shared_ptr> append_null_reader_; + std::shared_ptr state_; }; class CloseTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { public: - CloseTrackingRealtimeStoreFactory(const std::shared_ptr>& close_count, - const std::shared_ptr>& append_null_reader) - : close_count_(close_count), append_null_reader_(append_null_reader) {} + explicit CloseTrackingRealtimeStoreFactory( + const std::shared_ptr& state) + : state_(state) {} Result> Create(RealtimeStoreCreateRequest&& request) override { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, delegate_.Create(std::move(request))); - return std::shared_ptr(std::make_shared( - delegate, close_count_, append_null_reader_)); + return std::shared_ptr( + std::make_shared(delegate, state_)); } private: ArrowRealtimeStoreFactory delegate_; - std::shared_ptr> close_count_; - std::shared_ptr> append_null_reader_; + std::shared_ptr state_; }; class InvalidReaderRealtimeStore final : public RealtimeStore { @@ -1727,12 +1751,10 @@ TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestPkPluginQueryReaderCloseLifecycle) { +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderClose) { CreatePkTable(); - auto close_count = std::make_shared>(0); - auto append_null_reader = std::make_shared>(false); - auto factory = - std::make_shared(close_count, append_null_reader); + auto state = std::make_shared(); + auto factory = std::make_shared(state); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create(factory)); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -1758,15 +1780,71 @@ TEST_F(RealtimeWriteInteTest, TestPkPluginQueryReaderCloseLifecycle) { ASSERT_OK_AND_ASSIGN(std::unique_ptr explicitly_closed_reader, create_reader()); explicitly_closed_reader->Close(); explicitly_closed_reader.reset(); - ASSERT_EQ(1, close_count->load(std::memory_order_acquire)); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); ASSERT_OK_AND_ASSIGN(std::unique_ptr destroyed_reader, create_reader()); destroyed_reader.reset(); - ASSERT_EQ(2, close_count->load(std::memory_order_acquire)); + ASSERT_EQ(2, state->query_close_count->load(std::memory_order_acquire)); + + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch({Row{2, "two", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); - append_null_reader->store(true, std::memory_order_release); - ASSERT_NOK_WITH_MSG(create_reader(), "PK real-time store returned a null query reader"); - ASSERT_EQ(3, close_count->load(std::memory_order_acquire)); + auto create_reader = [&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + return table_read->CreateReader(plan->Splits()); + }; + + for (int32_t null_index = 0; null_index <= 2; ++null_index) { + state->query_null_index = null_index; + ASSERT_NOK_WITH_MSG(create_reader(), "PK real-time store returned a null query reader"); + ASSERT_EQ(2 * (null_index + 1), state->query_close_count->load(std::memory_order_acquire)); + } + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCommitReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + state->commit_null_index = 1; + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PK real-time store returned a null commit reader"); + ASSERT_EQ(1, state->commit_close_count->load(std::memory_order_acquire)); ASSERT_OK(writer->Close()); } From 2df7b78dfd6fd806e7a9f4baba4b2bc6a6c71680 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:05:54 +0800 Subject: [PATCH 07/24] fix(read): preserve PK reader metrics after close --- src/paimon/core/table/source/key_value_table_read.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 9b5f6ee83..76160ac93 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -85,6 +85,7 @@ class QueryBatchKeyValueReader final : public KeyValueRecordReader { std::shared_ptr row_kinds_; std::shared_ptr key_context_; std::shared_ptr value_context_; + bool closed_ = false; }; class QueryBatchKeyValueReader::Iterator final : public KeyValueRecordReader::Iterator { @@ -160,6 +161,10 @@ std::shared_ptr QueryBatchKeyValueReader::GetReaderMetrics() const { } void QueryBatchKeyValueReader::Close() { + if (closed_) { + return; + } + closed_ = true; values_.reset(); sequences_.reset(); row_kinds_.reset(); @@ -167,7 +172,6 @@ void QueryBatchKeyValueReader::Close() { value_context_.reset(); if (reader_) { reader_->Close(); - reader_.reset(); } } From 7148081bf4a675dabb313bdc730b4c8194a86f35 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:06:13 +0800 Subject: [PATCH 08/24] refactor(realtime): colocate PK realtime option validation --- .../core/operation/file_store_write.cpp | 3 +- .../realtime/primary_key_realtime_store.cpp | 34 +++++++++++++++++++ .../realtime/primary_key_realtime_store.h | 3 ++ .../primary_key_realtime_store_test.cpp | 26 ++++++++++++++ src/paimon/core/table/source/table_scan.cpp | 4 +-- .../core/utils/primary_key_table_utils.cpp | 32 ----------------- .../core/utils/primary_key_table_utils.h | 2 -- .../utils/primary_key_table_utils_test.cpp | 25 -------------- 8 files changed, 67 insertions(+), 62 deletions(-) diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index f216476bd..4d4f45156 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -36,6 +36,7 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" @@ -197,7 +198,7 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options)); if (ignore_previous_files) { return Status::NotImplemented( "PK realtime v1 requires restore from the latest snapshot"); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 84afb97a4..afdc0c73c 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -36,6 +36,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/core_options.h" #include "paimon/core/io/key_value_in_memory_record_reader.h" #include "paimon/core/io/key_value_projection_consumer.h" #include "paimon/core/io/key_value_projection_reader.h" @@ -45,6 +46,39 @@ #include "paimon/macros.h" namespace paimon { + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime v1 requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime v1 does not support data evolution"); + } + if (!options.GetFieldsSequenceGroups().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence groups"); + } + if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || + options.AggregationRemoveRecordOnDelete() || + !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { + return Status::NotImplemented("PK realtime v1 requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime v1 supports only ascending sequence.field.sort-order"); + } + if (options.NeedLookup() || options.DeletionVectorsEnabled() || + options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); + } + return Status::OK(); +} + namespace { uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h index 05225ed19..017864c04 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.h +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -33,6 +33,7 @@ class Schema; namespace paimon { +class CoreOptions; class FieldsComparator; struct KeyValue; class MemoryPool; @@ -40,6 +41,8 @@ class InternalRow; template class MergeFunctionWrapper; +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options); + /// Optional metadata exposed by PK query readers with a known inclusive key range. class PrimaryKeyRangeProvider { public: diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index 9da272e0f..cbbf9c82a 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -19,6 +19,7 @@ #include "paimon/core/realtime/primary_key_realtime_store.h" +#include #include #include #include @@ -30,6 +31,7 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/core_options.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/memory/memory_pool.h" @@ -37,6 +39,30 @@ namespace paimon::test { +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestSupportedOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options)); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options)); + } +} + class PrimaryKeyRealtimeStoreTest : public testing::Test { public: void SetUp() override { diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 92155de3b..dcf10e90c 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -41,6 +41,7 @@ #include "paimon/core/operation/data_evolution_file_store_scan.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/schema_validation.h" @@ -63,7 +64,6 @@ #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/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/format/file_format.h" #include "paimon/realtime/realtime_context.h" @@ -233,7 +233,7 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& return Status::Invalid("real-time union read does not support data evolution"); } if (!table_schema.PrimaryKeys().empty()) { - PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(core_options)); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options)); } if (context.IsStreamingMode()) { return Status::Invalid("real-time union read currently supports batch scans only"); diff --git a/src/paimon/core/utils/primary_key_table_utils.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index 823d48c41..cf72da4ae 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -96,36 +96,4 @@ Result> PrimaryKeyTableUtils::CreateSequenceFi options.SequenceFieldSortOrderIsAscending()); } -Status PrimaryKeyTableUtils::ValidateRealtimeOptions(const CoreOptions& options) { - if (options.GetBucket() <= 0) { - return Status::NotImplemented("PK realtime v1 requires fixed buckets"); - } - if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { - return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); - } - if (options.DataEvolutionEnabled()) { - return Status::NotImplemented("PK realtime v1 does not support data evolution"); - } - if (!options.GetFieldsSequenceGroups().empty()) { - return Status::NotImplemented("PK realtime v1 does not support sequence groups"); - } - if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || - options.AggregationRemoveRecordOnDelete() || - !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { - return Status::NotImplemented("PK realtime v1 requires default delete behavior"); - } - if (!options.GetSequenceField().empty()) { - return Status::NotImplemented("PK realtime v1 does not support sequence.field"); - } - if (!options.SequenceFieldSortOrderIsAscending()) { - return Status::NotImplemented( - "PK realtime v1 supports only ascending sequence.field.sort-order"); - } - if (options.NeedLookup() || options.DeletionVectorsEnabled() || - options.GetChangelogProducer() != ChangelogProducer::NONE) { - return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); - } - return Status::OK(); -} - } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index 7877ee4ab..c40e92cda 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -58,8 +58,6 @@ class PrimaryKeyTableUtils { static Result> CreateSequenceFieldsComparator( const std::vector& value_fields, const CoreOptions& options); - - static Status ValidateRealtimeOptions(const CoreOptions& options); }; } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 072965cff..1a7345fdf 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -19,7 +19,6 @@ #include "paimon/core/utils/primary_key_table_utils.h" #include -#include #include #include #include @@ -114,28 +113,4 @@ TEST(PrimaryKeyTableUtilsTest, TestCreateFirstRowMergeFunctionWithIgnoreDelete) "First row merge engine can not accept DELETE/UPDATE_BEFORE records"); } -TEST(PrimaryKeyTableUtilsTest, TestSupportedRealtimeOptions) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); - ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); -} - -TEST(PrimaryKeyTableUtilsTest, TestUnsupportedRealtimeOptions) { - const std::string sequence_group = - std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; - const std::vector> unsupported_options = { - {{Options::BUCKET, "0"}}, - {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, - {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, - {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, - {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, - {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, - {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, - {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, - }; - for (const std::map& option_map : unsupported_options) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); - ASSERT_NOK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options)); - } -} - } // namespace paimon::test From 8273045654eb71a78b675fc6be0ad0d0a85c69b6 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:45:39 +0800 Subject: [PATCH 09/24] test(realtime): improve primary key coverage --- .../primary_key_realtime_store_test.cpp | 281 ++++++++++++++---- test/inte/realtime_write_inte_test.cpp | 248 +++++++++++++++- 2 files changed, 463 insertions(+), 66 deletions(-) diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index cbbf9c82a..5c04d4310 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -19,9 +19,13 @@ #include "paimon/core/realtime/primary_key_realtime_store.h" +#include +#include #include #include +#include #include +#include #include #include "arrow/api.h" @@ -29,7 +33,6 @@ #include "arrow/ipc/json_simple.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" #include "paimon/core/core_options.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" @@ -69,24 +72,37 @@ class PrimaryKeyRealtimeStoreTest : public testing::Test { pool_ = std::shared_ptr(GetMemoryPool()); schema_ = arrow::schema( {arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); - ASSERT_OK_AND_ASSIGN(key_comparator_, - FieldsComparator::Create({DataField(0, schema_->field(0))}, - /*is_ascending_order=*/true)); + ASSERT_OK_AND_ASSIGN(store_, CreateStore(schema_, {"id"}, /*restore_max_sequence=*/4)); + } + + Result> CreateStore( + const std::shared_ptr& schema, const std::vector& primary_keys, + int64_t restore_max_sequence) const { + std::vector key_fields; + key_fields.reserve(primary_keys.size()); + for (const std::string& primary_key : primary_keys) { + const int32_t index = schema->GetFieldIndex(primary_key); + key_fields.emplace_back(index, schema->field(index)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, + /*is_ascending_order=*/true)); auto merge_factory = []() { auto merge_function = std::make_unique(/*ignore_delete=*/false); return std::make_shared(std::move(merge_function)); }; - ASSERT_OK_AND_ASSIGN( - store_, PrimaryKeyRealtimeStore::Create(schema_, {"id"}, key_comparator_, merge_factory, - /*restore_max_sequence_number=*/4, - /*read_batch_size=*/1024, pool_)); + return PrimaryKeyRealtimeStore::Create(schema, primary_keys, key_comparator, merge_factory, + restore_max_sequence, + /*read_batch_size=*/2, pool_); } std::unique_ptr MakeBatch( - const std::string& json, const std::vector& row_kinds = {}) const { + const std::string& json, const std::vector& row_kinds = {}, + const std::shared_ptr& schema = nullptr) const { + const std::shared_ptr& batch_schema = schema ? schema : schema_; std::shared_ptr array = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json) + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(batch_schema->fields()), json) .ValueOrDie(); ArrowArray c_array; EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); @@ -95,35 +111,39 @@ class PrimaryKeyRealtimeStoreTest : public testing::Test { return builder.Finish().value(); } - std::unique_ptr MakeReadSchema(bool include_sequence) const { - arrow::FieldVector fields; - if (include_sequence) { - fields.push_back( - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())); - } - fields.insert(fields.end(), schema_->fields().begin(), schema_->fields().end()); + std::unique_ptr MakeReadSchema(const arrow::FieldVector& fields) const { auto c_schema = std::make_unique(); EXPECT_TRUE(arrow::ExportSchema(*arrow::schema(fields), c_schema.get()).ok()); return c_schema; } - void AssertReaderOutput(BatchReader* reader, const std::shared_ptr& type, + void AssertReaderOutput(const std::vector>& readers, + const std::shared_ptr& type, const std::string& json) const { - ASSERT_NE(nullptr, reader); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); - ASSERT_FALSE(BatchReader::IsEofBatch(batch)); - arrow::Result> imported_result = - arrow::ImportArray(batch.first.get(), batch.second.get()); - ASSERT_TRUE(imported_result.ok()) << imported_result.status().ToString(); - std::shared_ptr actual = std::move(imported_result).ValueOrDie(); + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + arrow::Result> imported = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported.ok()) << imported.status().ToString(); + batches.push_back(std::move(imported).ValueOrDie()); + } + } + ASSERT_FALSE(batches.empty()); + arrow::Result> concatenated = arrow::Concatenate(batches); + ASSERT_TRUE(concatenated.ok()) << concatenated.status().ToString(); + std::shared_ptr actual = std::move(concatenated).ValueOrDie(); std::shared_ptr expected = arrow::ipc::internal::json::ArrayFromJSON(type, json).ValueOrDie(); ASSERT_TRUE(actual->Equals(*expected)) << "expected: " << expected->ToString() << ", actual: " << actual->ToString(); - - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, reader->NextBatch()); - ASSERT_TRUE(BatchReader::IsEofBatch(eof)); - reader->Close(); + for (const std::unique_ptr& reader : readers) { + reader->Close(); + } } std::shared_ptr CommitType() const { @@ -143,10 +163,18 @@ class PrimaryKeyRealtimeStoreTest : public testing::Test { }); } + arrow::FieldVector FullQueryFields( + const std::shared_ptr& schema = nullptr) const { + const std::shared_ptr& query_schema = schema ? schema : schema_; + arrow::FieldVector fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; + fields.insert(fields.end(), query_schema->fields().begin(), query_schema->fields().end()); + return fields; + } + protected: std::shared_ptr pool_; std::shared_ptr schema_; - std::shared_ptr key_comparator_; std::shared_ptr store_; }; @@ -172,30 +200,48 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestWriteAndSeal) { ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); ASSERT_GT(store_->GetMemoryUsage(), 0); - auto merge_factory = []() { - auto merge_function = std::make_unique(/*ignore_delete=*/false); - return std::make_shared(std::move(merge_function)); + struct ValidationCase { + int64_t restore_max_sequence; + std::string error; + }; + const std::vector cases = { + {-2, "restore max sequence number must be at least -1"}, + {std::numeric_limits::max(), "sequence number has reached INT64_MAX"}, }; - ASSERT_NOK_WITH_MSG(PrimaryKeyRealtimeStore::Create( - schema_, {"id"}, key_comparator_, merge_factory, - /*restore_max_sequence_number=*/-2, /*read_batch_size=*/1024, pool_), - "restore max sequence number must be at least -1"); + for (const ValidationCase& test_case : cases) { + ASSERT_NOK_WITH_MSG(CreateStore(schema_, {"id"}, test_case.restore_max_sequence), + test_case.error); + } } -TEST_F(PrimaryKeyRealtimeStoreTest, TestCommitReaderPreservesMutations) { +TEST_F(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { + ASSERT_OK(store_->Write(RealtimeWriteBatch{ + MakeBatch(R"([[3, "three"], [1, "before"]])", + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_BEFORE}), + OffsetRange(0, 2)})); + ASSERT_OK(store_->Write(RealtimeWriteBatch{ + MakeBatch(R"([[2, "after"]])", {RecordBatch::RowKind::UPDATE_AFTER}), OffsetRange(2, 3)})); ASSERT_OK(store_->Write( - RealtimeWriteBatch{MakeBatch(R"([[2, "old"], [1, "one"], [2, "new"]])", - {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, - RecordBatch::RowKind::UPDATE_AFTER}), - OffsetRange(0, 3)})); + RealtimeWriteBatch{MakeBatch(R"([[4, "deleted"], [0, "zero"]])", + {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}), + OffsetRange(3, 5)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, store_->SealForCommit()); ASSERT_TRUE(segment.has_value()); ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateCommitReaders(segment.value())); - ASSERT_EQ(1, readers.size()); - AssertReaderOutput(readers[0].get(), CommitType(), - R"([[0, 2, "old"], [0, 1, "one"], [2, 2, "new"]])"); + AssertReaderOutput(readers, CommitType(), + R"([[0, 3, "three"], [1, 1, "before"], [2, 2, "after"], + [3, 4, "deleted"], [0, 0, "zero"]])"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + AssertReaderOutput(readers, QueryType(), + R"([[0, 9, 0, "zero"], [1, 6, 1, "before"], [2, 7, 2, "after"], + [0, 5, 3, "three"], [3, 8, 4, "deleted"]])"); } TEST_F(PrimaryKeyRealtimeStoreTest, TestMutationMerge) { @@ -207,13 +253,12 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestMutationMerge) { OffsetRange(2, 4)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - ASSERT_EQ(1, readers.size()); - AssertReaderOutput(readers[0].get(), QueryType(), R"([[2, 7, 1, "new"], [3, 8, 2, "gone"]])"); + AssertReaderOutput(readers, QueryType(), R"([[2, 7, 1, "new"], [3, 8, 2, "gone"]])"); } TEST_F(PrimaryKeyRealtimeStoreTest, TestReadViewLifecycle) { @@ -230,15 +275,14 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestReadViewLifecycle) { ASSERT_OK( store_->Write(RealtimeWriteBatch{MakeBatch(R"([[13, "later"]])"), OffsetRange(13, 14)})); - std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, /*offset_begin=*/11, context)); - ASSERT_EQ(1, readers.size()); - AssertReaderOutput(readers[0].get(), QueryType(), R"([[0, 6, 11, "b"], [0, 7, 12, "c"]])"); + AssertReaderOutput(readers, QueryType(), R"([[0, 6, 11, "b"], [0, 7, 12, "c"]])"); - std::unique_ptr empty_schema = MakeReadSchema(/*include_sequence=*/true); + std::unique_ptr empty_schema = MakeReadSchema(FullQueryFields()); context.read_schema = empty_schema.get(); ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/13, context)); ASSERT_TRUE(readers.empty()); @@ -251,20 +295,135 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryKeyRange) { ASSERT_OK(store_->Write( RealtimeWriteBatch{MakeBatch(R"([[9, "nine"], [7, "seven"]])"), OffsetRange(2, 4)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(/*include_sequence=*/true); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); ASSERT_EQ(2, readers.size()); - auto* first_range = dynamic_cast(readers[0].get()); - auto* second_range = dynamic_cast(readers[1].get()); - ASSERT_NE(nullptr, first_range); - ASSERT_NE(nullptr, second_range); - ASSERT_EQ(1, first_range->GetMinKey()->GetLong(0)); - ASSERT_EQ(5, first_range->GetMaxKey()->GetLong(0)); - ASSERT_EQ(7, second_range->GetMinKey()->GetLong(0)); - ASSERT_EQ(9, second_range->GetMaxKey()->GetLong(0)); + const std::vector> key_ranges = {{1, 5}, {7, 9}}; + for (size_t i = 0; i < readers.size(); ++i) { + auto* range = dynamic_cast(readers[i].get()); + ASSERT_NE(nullptr, range); + ASSERT_EQ(key_ranges[i].first, range->GetMinKey()->GetLong(0)); + ASSERT_EQ(key_ranges[i].second, range->GetMaxKey()->GetLong(0)); + } + AssertReaderOutput(readers, QueryType(), + R"([[0, 6, 1, "one"], [0, 5, 5, "five"], [0, 8, 7, "seven"], + [0, 7, 9, "nine"]])"); + + ASSERT_OK(store_->AdvanceCommittedOffset(2)); + ASSERT_OK_AND_ASSIGN(view, store_->AcquireReadView()); + read_schema = MakeReadSchema(FullQueryFields()); + context.read_schema = read_schema.get(); + ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + auto* range = dynamic_cast(readers[0].get()); + ASSERT_NE(nullptr, range); + ASSERT_EQ(7, range->GetMinKey()->GetLong(0)); + ASSERT_EQ(9, range->GetMaxKey()->GetLong(0)); + AssertReaderOutput(readers, QueryType(), R"([[0, 8, 7, "seven"], [0, 7, 9, "nine"]])"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestSequenceExhaustion) { + const int64_t max_sequence = std::numeric_limits::max(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(schema_, {"id"}, max_sequence - 3)); + ASSERT_OK(store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "kept"]])"), OffsetRange(10, 11)})); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[7, "rejected-a"], [8, "rejected-b"], [9, "rejected-c"]])"), + OffsetRange(11, 14)}), + "sequence range exceeds INT64_MAX"); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[2, "also-kept"]])"), OffsetRange(11, 12)})); + + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(10, 12), segment.value()->GetOffsetRange()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + ASSERT_EQ(std::optional(OffsetRange(10, 12)), view->GetOffsetRange()); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/10, context)); + AssertReaderOutput(readers, QueryType(), + R"([[0, 9223372036854775805, 1, "kept"], + [0, 9223372036854775806, 2, "also-kept"]])"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjection) { + ASSERT_OK( + store_->Write(RealtimeWriteBatch{MakeBatch(R"([[2, "b"], [1, "a"]])"), OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); + const std::shared_ptr value_kind = + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()); + const std::shared_ptr sequence = + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()); + struct ProjectionCase { + arrow::FieldVector requested; + std::shared_ptr expected_type; + std::string expected_json; + }; + const std::vector cases = { + {{schema_->field(1), value_kind, sequence, schema_->field(0)}, + arrow::struct_({value_kind, schema_->field(1), sequence, schema_->field(0)}), + R"([[0, "a", 6, 1], [0, "b", 5, 2]])"}, + {{schema_->field(0), value_kind}, + arrow::struct_({value_kind, schema_->field(0)}), + R"([[0, 1], [0, 2]])"}, + }; + for (const ProjectionCase& test_case : cases) { + std::unique_ptr read_schema = MakeReadSchema(test_case.requested); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); + AssertReaderOutput(readers, test_case.expected_type, test_case.expected_json); + } + + std::unique_ptr read_schema = + MakeReadSchema({arrow::field("unknown", arrow::int64())}); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_NOK_WITH_MSG(store_->CreateQueryReaders(view, /*offset_begin=*/0, context), + "query field is missing from write schema: unknown"); +} + +TEST_F(PrimaryKeyRealtimeStoreTest, TestCompositeKeyClipping) { + std::shared_ptr composite_schema = + arrow::schema({arrow::field("id", arrow::int64()), arrow::field("region", arrow::utf8()), + arrow::field("value", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(composite_schema, {"id", "region"}, + /*restore_max_sequence=*/4)); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[9, "z", "clipped"], [2, "b", "two-b"], [1, "c", "one-c"], + [2, "a", "two-a"]])", + {}, composite_schema), + OffsetRange(20, 24)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + std::unique_ptr read_schema = MakeReadSchema(FullQueryFields(composite_schema)); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/21, context)); + ASSERT_EQ(1, readers.size()); + auto* range = dynamic_cast(readers[0].get()); + ASSERT_NE(nullptr, range); + ASSERT_EQ(1, range->GetMinKey()->GetLong(0)); + ASSERT_EQ("c", range->GetMinKey()->GetString(1).ToString()); + ASSERT_EQ(2, range->GetMaxKey()->GetLong(0)); + ASSERT_EQ("b", range->GetMaxKey()->GetString(1).ToString()); + std::shared_ptr query_type = arrow::struct_( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()), + composite_schema->field(0), composite_schema->field(1), composite_schema->field(2)}); + AssertReaderOutput(readers, query_type, + R"([[0, 7, 1, "c", "one-c"], [0, 8, 2, "a", "two-a"], + [0, 6, 2, "b", "two-b"]])"); } } // namespace paimon::test diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index e6000561f..aad9dc2ac 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -646,16 +646,18 @@ class RealtimeWriteInteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } - void CreatePkTable(const std::vector& partition_keys = {}) const { + void CreatePkTable(const std::vector& partition_keys = {}, + const std::vector& primary_keys = {"id"}) const { auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*schema_, c_schema.get()).ok()); ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, Catalog::Create(dir_->Str(), options_)); ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false)); - std::vector primary_keys = partition_keys; - primary_keys.push_back("id"); + std::vector table_primary_keys = partition_keys; + table_primary_keys.insert(table_primary_keys.end(), primary_keys.begin(), + primary_keys.end()); ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), c_schema.get(), partition_keys, - primary_keys, options_, /*ignore_if_exists=*/false)); + table_primary_keys, options_, /*ignore_if_exists=*/false)); } Result> CreateRealtimeWriter( @@ -692,7 +694,7 @@ class RealtimeWriteInteTest : public ::testing::Test { std::string json = "["; for (size_t i = 0; i < rows.size(); ++i) { const auto& [id, payload, pt] = rows[i]; - if (pt != partition) { + if (partitioned && pt != partition) { return Status::Invalid("one test batch must contain only one partition"); } if (i > 0) { @@ -1383,6 +1385,242 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { ASSERT_TRUE(query_view->expired()); } +TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { + options_[Options::READ_BATCH_SIZE] = "2"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}, {3, "disk-3", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + }; + int64_t commit_identifier = 0; + for (const std::vector& disk_rows : disk_batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ++commit_identifier; + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "sealed-1", "p0"}, Row{2, "deleted-2", "p0"}, Row{4, "sealed-4", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "active-1", "p0"}, Row{4, "deleted-4", "p0"}, Row{5, "active-5", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"payload", "id"}, /*predicate=*/nullptr, + /*enable_predicate_filter=*/false)); + ASSERT_NE(nullptr, result.data); + ASSERT_GT(result.data->num_chunks(), 1); + for (const std::shared_ptr& chunk : result.data->chunks()) { + ASSERT_LE(chunk->length(), 2); + } + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()), + arrow::field("id", arrow::int64())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "active-1", 1], + [0, "disk-3", 3], + [0, "active-5", 5], + [0, "disk-10", 10], + [0, "disk-11", 11] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { + CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({Row{1, "a", "disk-1a"}, Row{1, "b", "disk-1b"}, + Row{2, "a", "disk-2a"}, Row{3, "c", "disk-3c"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, disk_progress.size()); + ASSERT_EQ(OffsetRange(0, 4), disk_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(disk_progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "a", "sealed-1a"}, Row{1, "b", "deleted-1b"}, Row{2, "b", "sealed-2b"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + ASSERT_EQ(OffsetRange(4, 7), sealed_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(sealed_progress).size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "a", "active-1a"}, Row{1, "c", "active-1c"}, Row{2, "a", "active-2a"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + ASSERT_FALSE(split->DiskSplits().empty()); + ASSERT_EQ(4, split->CommittedEndOffset()); + ASSERT_EQ(10, split->MemoryEndOffset()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "a", "active-1a"}, + {1, "c", "active-1c"}, + {2, "a", "active-2a"}, + {2, "b", "sealed-2b"}, + {3, "c", "disk-3c"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkPartitionBucketRecovery) { + options_[Options::BUCKET] = "2"; + CreatePkTable(/*partition_keys=*/{"pt"}); + const RealtimePartitionBucket p0b0({{"pt", "p0"}}, /*bucket=*/0); + const RealtimePartitionBucket p1b1({{"pt", "p1"}}, /*bucket=*/1); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_first_batch, + MakeBatch({Row{0, "p0-zero", "p0"}, Row{1, "p0-one", "p0"}}, + /*partitioned=*/true, /*bucket=*/0)); + ASSERT_OK(first_writer->Write(std::move(p0_first_batch))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p1_first_batch, + MakeBatch({Row{10, "p1-ten", "p1"}, Row{11, "p1-eleven", "p1"}, Row{12, "p1-twelve", "p1"}}, + /*partitioned=*/true, /*bucket=*/1)); + ASSERT_OK(first_writer->Write(std::move(p1_first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, first_progress.size()); + std::map first_ranges; + std::map> first_sequences; + for (const RealtimeCommitProgress& progress : first_progress) { + first_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + first_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(0, 2), first_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(0, 3), first_ranges.at(p1b1)); + ASSERT_EQ((std::make_pair(0, 1)), first_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(0, 2)), first_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, + Commit(first_progress, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->RefreshCommittedSnapshot(first_snapshot_id)); + ASSERT_OK(first_writer->Close()); + first_writer.reset(); + first_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p0_second_batch, + MakeBatch({Row{0, "p0-zero-new", "p0"}, Row{2, "p0-two", "p0"}}, + /*partitioned=*/true, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p0_second_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_second_batch, + MakeBatch({Row{10, "p1-ten-deleted", "p1"}, Row{13, "p1-thirteen", "p1"}}, + /*partitioned=*/true, /*bucket=*/1, + {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p1_second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, second_progress.size()); + std::map second_ranges; + std::map> second_sequences; + for (const RealtimeCommitProgress& progress : second_progress) { + second_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + second_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(2, 4), second_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(3, 5), second_ranges.at(p1b1)); + ASSERT_EQ((std::make_pair(2, 3)), second_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(3, 4)), second_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_progress, /*commit_identifier=*/1)); + ASSERT_OK(second_writer->RefreshCommittedSnapshot(second_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(second_context)); + std::sort(actual_rows.begin(), actual_rows.end()); + ASSERT_EQ((std::vector{{0, "p0-zero-new", "p0"}, + {1, "p0-one", "p0"}, + {2, "p0-two", "p0"}, + {11, "p1-eleven", "p1"}, + {12, "p1-twelve", "p1"}, + {13, "p1-thirteen", "p1"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); + + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(2, offsets.size()); + ASSERT_EQ(4, offsets.at(p0b0)); + ASSERT_EQ(5, offsets.at(p1b1)); +} + TEST_F(RealtimeWriteInteTest, TestPkRecovery) { CreatePkTable(); From 728b97ed00fd4d11333e7224b03108296fffe94f Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:38:57 +0800 Subject: [PATCH 10/24] fix(realtime): prevent sequence reuse and align nested projections --- .../realtime/primary_key_realtime_store.cpp | 9 +++ .../primary_key_realtime_store_test.cpp | 28 +++++++ .../core/realtime/realtime_context_impl.cpp | 10 ++- .../core/realtime/realtime_context_test.cpp | 16 +++- test/inte/realtime_write_inte_test.cpp | 81 +++++++++++++++++++ 5 files changed, 139 insertions(+), 5 deletions(-) diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index afdc0c73c..7999de75d 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -43,6 +43,7 @@ #include "paimon/core/io/merged_key_value_record_reader.h" #include "paimon/core/key_value.h" #include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h" +#include "paimon/core/utils/nested_projection_utils.h" #include "paimon/macros.h" namespace paimon { @@ -411,6 +412,7 @@ class PrimaryKeyRealtimeStore::Impl { arrow::ImportSchema(context.read_schema)); arrow::FieldVector output_fields = { DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; + arrow::FieldVector aligned_value_fields = write_schema_->fields(); std::vector projection = {KeyValueProjectionConsumer::kValueKindProjection}; for (const std::shared_ptr& field : requested->fields()) { if (field->name() == SpecialFields::ValueKind().Name()) { @@ -426,8 +428,11 @@ class PrimaryKeyRealtimeStore::Impl { return Status::Invalid("PK real-time query field is missing from write schema: ", field->name()); } + aligned_value_fields[index] = field; projection.push_back(index); } + const std::shared_ptr aligned_value_type = + arrow::struct_(aligned_value_fields); std::vector> result; for (const BatchGroup& group : typed->Groups()) { @@ -452,6 +457,10 @@ class PrimaryKeyRealtimeStore::Impl { if (!max_key || key_comparator_->CompareTo(*key_range.second, *max_key) > 0) { max_key = key_range.second; } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr aligned, + NestedProjectionUtils::AlignArrayToReadType( + selected, aligned_value_type, arrow_pool_.get())); + selected = checked_pointer_cast(aligned); std::vector selected_kinds; if (!batch->row_kinds.empty()) { selected_kinds.assign(batch->row_kinds.begin() + offset, diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index 5c04d4310..ef293e54b 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -392,6 +392,34 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjection) { "query field is missing from write schema: unknown"); } +TEST_F(PrimaryKeyRealtimeStoreTest, TestNestedProjection) { + const std::shared_ptr id = + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))); + const std::shared_ptr a = + DataField::ConvertDataFieldToArrowField(DataField(10, arrow::field("a", arrow::int64()))); + const std::shared_ptr b = + DataField::ConvertDataFieldToArrowField(DataField(11, arrow::field("b", arrow::int64()))); + const std::shared_ptr payload = DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("payload", arrow::struct_({a, b})))); + const std::shared_ptr nested_schema = arrow::schema({id, payload}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(nested_schema, {"id"}, /*restore_max_sequence=*/4)); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[2, [200, 2000]], [1, [100, null]], [3, [300, 3000]]])", {}, nested_schema), + OffsetRange(0, 3)})); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + const std::shared_ptr projected_payload = payload->WithType(arrow::struct_({b})); + std::unique_ptr read_schema = MakeReadSchema({projected_payload}); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + const std::shared_ptr result_type = arrow::struct_( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), projected_payload}); + AssertReaderOutput(readers, result_type, R"([[0, [null]], [0, [2000]], [0, [3000]]])"); +} + TEST_F(PrimaryKeyRealtimeStoreTest, TestCompositeKeyClipping) { std::shared_ptr composite_schema = arrow::schema({arrow::field("id", arrow::int64()), arrow::field("region", arrow::utf8()), diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 6624059a6..066e54e8a 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -82,6 +82,7 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); const RealtimePartitionBucket key(request.partition, request.bucket); + auto iter = stores_.find(key); std::optional initial_max_sequence_number; PrimaryKeyRealtimeStoreCreateConfig* primary_key_config = std::get_if(&request.mode_config); @@ -89,6 +90,14 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( auto [sequence_iter, inserted] = materialized_max_sequence_numbers_.emplace( key, primary_key_config->restore_max_sequence_number); if (!inserted && primary_key_config->restore_max_sequence_number > sequence_iter->second) { + if (iter != stores_.end()) { + if (request.write_schema) { + ArrowSchemaRelease(request.write_schema.get()); + } + return Status::Invalid( + "restore max sequence number exceeds the materialized watermark of an " + "existing PK real-time store"); + } sequence_iter->second = primary_key_config->restore_max_sequence_number; } initial_max_sequence_number = sequence_iter->second; @@ -105,7 +114,6 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( } initial_offset = offset_iter->second; } - auto iter = stores_.find(key); if (iter != stores_.end()) { if (request.write_schema) { ArrowSchemaRelease(request.write_schema.get()); diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index b4d2c6718..ab0abe4a7 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -198,12 +198,20 @@ TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { ASSERT_EQ(first_state.store, retained_state.store); ASSERT_EQ(8, retained_state.initial_max_sequence_number); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState restored_state, + ASSERT_NOK_WITH_MSG( GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, + /*restore_max_sequence_number=*/10, GetDefaultPool()), + "restore max sequence number exceeds the materialized watermark of an " + "existing PK real-time store"); + + const RealtimePartitionBucket new_partition_bucket(partition, /*bucket=*/1); + context->AdvanceMaterializedMaxSequenceNumber(new_partition_bucket, + /*max_sequence_number=*/8); + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState new_state, + GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/1, /*restore_max_sequence_number=*/10, GetDefaultPool())); - ASSERT_EQ(first_state.store, restored_state.store); - ASSERT_EQ(10, restored_state.initial_max_sequence_number); + ASSERT_EQ(10, new_state.initial_max_sequence_number); } TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index aad9dc2ac..9f302eb37 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1458,6 +1458,87 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { + const std::shared_ptr projected_b = arrow::field("b", arrow::int64()); + fields_ = { + arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::struct_({arrow::field("a", arrow::int64()), projected_b})), + arrow::field("pt", arrow::utf8()), + }; + schema_ = arrow::schema(fields_); + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + auto make_batch = [&](const std::string& json) -> Result> { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetBucket(0).Finish(); + }; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + make_batch(R"([[1, [101, 1001], "p0"], [2, [102, 1002], "p0"]])")); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr sealed_batch, + make_batch(R"([[1, [201, 2001], "p0"], [3, [203, 2003], "p0"]])")); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr active_batch, + make_batch(R"([[1, [301, 3001], "p0"], [4, [304, null], "p0"]])")); + ASSERT_OK(writer->Write(std::move(active_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + auto projected_schema = arrow::schema({ + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadSchema(std::move(c_schema)) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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 reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(reader.get())); + const std::shared_ptr result_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + const std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, [3001], 1], + [0, [1002], 2], + [0, [2003], 3], + [0, [null], 4] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) + << actual->ToString(); + reader->Close(); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, From f3df0e37feaf32713d3feea59a023be386b5f702 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:20:03 +0800 Subject: [PATCH 11/24] fix(realtime): align PK reads across schema changes --- .../realtime/primary_key_realtime_store.cpp | 28 +- test/inte/CMakeLists.txt | 7 + ...chema_evolution_write_verify_inte_test.cpp | 1110 +++++++++++++++++ 3 files changed, 1136 insertions(+), 9 deletions(-) create mode 100644 test/inte/schema_evolution_write_verify_inte_test.cpp diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 7999de75d..6565ed8d7 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -423,12 +423,18 @@ class PrimaryKeyRealtimeStore::Impl { projection.push_back(KeyValueProjectionConsumer::kSequenceNumberProjection); continue; } - const int32_t index = write_schema_->GetFieldIndex(field->name()); + int32_t index = write_schema_->GetFieldIndex(field->name()); if (index < 0) { - return Status::Invalid("PK real-time query field is missing from write schema: ", - field->name()); + Result field_id = NestedProjectionUtils::GetPaimonFieldId(field); + if (!field_id.ok()) { + return Status::Invalid( + "PK real-time query field is missing from write schema: ", field->name()); + } + index = static_cast(aligned_value_fields.size()); + aligned_value_fields.push_back(field); + } else { + aligned_value_fields[index] = field; } - aligned_value_fields[index] = field; projection.push_back(index); } const std::shared_ptr aligned_value_type = @@ -446,8 +452,16 @@ class PrimaryKeyRealtimeStore::Impl { const int64_t offset = std::max(0, lower - batch->offset_range.begin); const int64_t length = batch->data->length() - offset; std::shared_ptr sliced = batch->data->Slice(offset, length); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr aligned, + NestedProjectionUtils::AlignArrayToReadType( + sliced, aligned_value_type, arrow_pool_.get())); + if (!aligned || aligned->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + "PK real-time query projection did not produce a " + "StructArray"); + } std::shared_ptr selected = - checked_pointer_cast(sliced); + checked_pointer_cast(aligned); using KeyRange = std::pair, std::shared_ptr>; PAIMON_ASSIGN_OR_RAISE(KeyRange key_range, GetKeyRange(selected)); @@ -457,10 +471,6 @@ class PrimaryKeyRealtimeStore::Impl { if (!max_key || key_comparator_->CompareTo(*key_range.second, *max_key) > 0) { max_key = key_range.second; } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr aligned, - NestedProjectionUtils::AlignArrayToReadType( - selected, aligned_value_type, arrow_pool_.get())); - selected = checked_pointer_cast(aligned); std::vector selected_kinds; if (!batch->row_kinds.empty()) { selected_kinds.assign(batch->row_kinds.begin() + offset, diff --git a/test/inte/CMakeLists.txt b/test/inte/CMakeLists.txt index 75147ce60..f1b3f8ce6 100644 --- a/test/inte/CMakeLists.txt +++ b/test/inte/CMakeLists.txt @@ -43,6 +43,13 @@ if(PAIMON_BUILD_TESTS) test_utils_static ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(schema_evolution_write_verify_inte_test + STATIC_LINK_LIBS + paimon_shared + ${TEST_STATIC_LINK_LIBS} + test_utils_static + ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(global_index_test STATIC_LINK_LIBS paimon_shared diff --git a/test/inte/schema_evolution_write_verify_inte_test.cpp b/test/inte/schema_evolution_write_verify_inte_test.cpp new file mode 100644 index 000000000..dcadbd9e1 --- /dev/null +++ b/test/inte/schema_evolution_write_verify_inte_test.cpp @@ -0,0 +1,1110 @@ +/* + * 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 +#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/common/utils/scope_guard.h" +#include "paimon/core/core_options.h" +#include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/snapshot.h" +#include "paimon/core/table/sink/commit_message_impl.h" +#include "paimon/core/table/source/data_split_impl.h" +#include "paimon/defs.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/file_index/file_index_reader.h" +#include "paimon/file_index/file_index_result.h" +#include "paimon/file_store_commit.h" +#include "paimon/file_store_write.h" +#include "paimon/fs/file_system.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/record_batch.h" +#include "paimon/scan_context.h" +#include "paimon/table/source/plan.h" +#include "paimon/table/source/startup_mode.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/test_helper.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/write_context.h" + +namespace paimon::test { +namespace { + +std::map BaseOptions() { + return {{Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, + {Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, {Options::TARGET_FILE_SIZE, "1MB"}}; +} + +std::map DataEvolutionOptions() { + return {{Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, + {Options::FILE_SYSTEM, "local"}, {Options::TARGET_FILE_SIZE, "1MB"}, + {Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}; +} + +arrow::FieldVector BaseFields() { + return {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8())}; +} + +arrow::FieldVector EvolvedFields() { + return {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), + arrow::field("extra", arrow::int32())}; +} + +arrow::FieldVector DataEvolutionFields() { + return {arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::utf8()), + arrow::field("f2", arrow::utf8())}; +} + +Result> MakeBatch( + const arrow::FieldVector& fields, const std::string& json, + const std::map& partition, int32_t bucket, + const std::vector& row_kinds = {}) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetPartition(partition).SetBucket(bucket).SetRowKinds(row_kinds).Finish(); +} + +Result> MakeUnbucketedBatch( + const arrow::FieldVector& fields, const std::string& json, + const std::map& partition, + const std::vector& row_kinds = {}) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetPartition(partition).SetRowKinds(row_kinds).Finish(); +} + +Result> CreateWriter( + const std::string& table_path, const std::map& options, + const std::shared_ptr& realtime_context = nullptr, + const std::vector& write_schema = {}) { + WriteContextBuilder builder(table_path, "schema_evolution_verify"); + builder.SetOptions(options).WithStreamingMode(true); + if (realtime_context) { + builder.WithRealtimeContext(realtime_context); + } + if (!write_schema.empty()) { + builder.WithWriteSchema(write_schema); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + return FileStoreWrite::Create(std::move(context)); +} + +Result>> WriteWithNewWriter( + const std::string& table_path, const std::map& options, + std::unique_ptr batch, int64_t commit_identifier, + const std::vector& write_schema = {}) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + CreateWriter(table_path, options, nullptr, write_schema)); + 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> CreateCommit( + const std::string& table_path, const std::map& options) { + CommitContextBuilder builder(table_path, "schema_evolution_verify"); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options).IgnoreEmptyCommit(false).Finish()); + return FileStoreCommit::Create(std::move(context)); +} + +Status CommitMessages(const std::string& table_path, + const std::map& options, + const std::vector>& messages, + int64_t commit_identifier) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + CreateCommit(table_path, options)); + return commit->Commit(messages, commit_identifier); +} + +Result CommitRealtimeMessages(const std::string& table_path, + const std::map& options, + const std::vector& messages, + int64_t commit_identifier) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + CreateCommit(table_path, options)); + return commit->CommitWithProgress(messages, commit_identifier, /*watermark=*/std::nullopt); +} + +Result> LatestSnapshot(const std::string& table_path, + const std::map& options, + const std::shared_ptr& file_system) { + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options, file_system)); + SnapshotManager snapshot_manager(core_options.GetFileSystem(), table_path); + return snapshot_manager.LatestSnapshot(); +} + +Result> ScanTable( + const std::string& table_path, const std::map& options, + const std::shared_ptr& pool, + const std::shared_ptr& realtime_context = nullptr, + const std::shared_ptr& predicate = nullptr) { + ScanContextBuilder scan_builder(table_path); + scan_builder.SetOptions(options) + .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()) + .SetPredicate(predicate) + .WithMemoryPool(pool); + if (realtime_context) { + scan_builder.WithRealtimeContext(realtime_context); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, scan_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_scan, + TableScan::Create(std::move(scan_context))); + return table_scan->CreatePlan(); +} + +std::vector> DataFilesFromPlan(const std::shared_ptr& plan) { + std::vector> files; + for (const std::shared_ptr& split : plan->Splits()) { + std::shared_ptr data_split = split; + if (std::shared_ptr indexed_split = + std::dynamic_pointer_cast(split)) { + data_split = indexed_split->GetDataSplit(); + } + std::shared_ptr split_impl = + std::dynamic_pointer_cast(data_split); + if (!split_impl) { + continue; + } + const std::vector>& split_files = split_impl->DataFiles(); + files.insert(files.end(), split_files.begin(), split_files.end()); + } + return files; +} + +size_t CountIndexedSplits(const std::shared_ptr& plan) { + size_t count = 0; + for (const std::shared_ptr& split : plan->Splits()) { + if (std::dynamic_pointer_cast(split)) { + count++; + } + } + return count; +} + +Status EvolveSchema(const std::string& table_path, const std::shared_ptr& file_system, + const std::vector& fields, int32_t highest_field_id, + const std::map& options) { + return TestHelper::WriteNextSchema(file_system, table_path, fields, highest_field_id, options); +} + +void AssignFirstRowId(const std::vector>& messages, + int64_t first_row_id) { + for (const std::shared_ptr& commit_message : messages) { + std::shared_ptr message = + std::dynamic_pointer_cast(commit_message); + ASSERT_TRUE(message); + for (const std::shared_ptr& file : + message->GetNewFilesIncrement().NewFiles()) { + file->AssignFirstRowId(first_row_id); + } + } +} + +struct CollectedReadResult { + std::unique_ptr table_read; + std::unique_ptr reader; + std::shared_ptr data; +}; + +Result ReadRows( + const std::string& table_path, const std::map& options, + const std::shared_ptr& pool, + const std::shared_ptr& realtime_context = nullptr, + const std::shared_ptr& predicate = nullptr, bool enable_predicate_filter = true) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + ScanTable(table_path, options, pool, realtime_context, predicate)); + + ReadContextBuilder read_builder(table_path); + read_builder.SetOptions(options) + .SetPredicate(predicate) + .EnablePredicateFilter(enable_predicate_filter) + .WithMemoryPool(pool); + if (realtime_context) { + read_builder.WithRealtimeContext(realtime_context); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, + table_read->CreateReader(plan->Splits())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr rows, + ReadResultCollector::CollectResult(batch_reader.get())); + return CollectedReadResult{std::move(table_read), std::move(batch_reader), std::move(rows)}; +} + +void AssertResultEquals(const std::shared_ptr& actual, + const arrow::FieldVector& fields, const std::string& expected_json) { + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + std::shared_ptr expected_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_with_row_kind), + expected_json) + .ValueOrDie(); + auto expected = std::make_shared(expected_array); + ASSERT_TRUE(expected->Equals(actual, arrow::EqualOptions::Defaults().diff_sink(&std::cout))) + << actual->ToString(); +} + +Status CreateTable(const std::string& warehouse, const std::shared_ptr& schema, + const std::vector& primary_keys, + const std::map& options) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr catalog, Catalog::Create(warehouse, options)); + PAIMON_RETURN_NOT_OK(catalog->CreateDatabase("foo", options, /*ignore_if_exists=*/false)); + ArrowSchema c_schema; + ArrowSchemaMarkReleased(&c_schema); + ScopeGuard guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); + return catalog->CreateTable(Identifier("foo", "bar"), &c_schema, + /*partition_keys=*/{}, primary_keys, options, + /*ignore_if_exists=*/false); +} + +Result> CreateFileIndexReader( + const std::shared_ptr& data_file, const std::shared_ptr& pool) { + if (data_file->embedded_index == nullptr) { + return Status::Invalid("data file does not contain an embedded file index"); + } + auto input = std::make_shared(data_file->embedded_index->data(), + data_file->embedded_index->size()); + return FileIndexFormat::CreateReader(input, pool); +} + +Result>> ReadEmbeddedIndexColumn( + const std::shared_ptr& data_file, const std::shared_ptr& schema, + const std::string& column, const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + CreateFileIndexReader(data_file, pool)); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, c_schema.get())); + return reader->ReadColumnIndex(column, c_schema.get()); +} + +class SchemaEvolutionWriteVerifyTest : public ::testing::Test { + protected: + void SetUp() override { + pool_ = GetDefaultPool(); + dir_ = UniqueTestDirectory::Create("local"); + ASSERT_TRUE(dir_); + table_path_ = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + } + + void TearDown() override { + dir_.reset(); + } + + std::shared_ptr pool_; + std::unique_ptr dir_; + std::string table_path_; +}; + +TEST_F(SchemaEvolutionWriteVerifyTest, + NonRealtimeAppendOldWriterCommitsOldSchemaFileIntoNewSchemaSnapshot) { + std::map options = BaseOptions(); + options["file-index.bitmap.columns"] = "payload"; + options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"], [2, "skip"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> messages, + old_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/1)); + ASSERT_EQ(1, messages.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(messages[0]); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + std::shared_ptr old_file = message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(0, old_file->schema_id); + ASSERT_TRUE(old_file->embedded_index); + ASSERT_TRUE(old_file->extra_files.empty()); + ASSERT_OK_AND_ASSIGN( + std::vector> payload_indexes, + ReadEmbeddedIndexColumn(old_file, arrow::schema(BaseFields()), "payload", pool_)); + ASSERT_EQ(1, payload_indexes.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr payload_hit, + payload_indexes[0]->VisitEqual(Literal(FieldType::STRING, "old", 3))); + ASSERT_OK_AND_ASSIGN(bool payload_remain, payload_hit->IsRemain()); + ASSERT_TRUE(payload_remain); + + ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::optional snapshot, + LatestSnapshot(table_path_, options, dir_->GetFileSystem())); + ASSERT_TRUE(snapshot); + ASSERT_EQ(1, snapshot->SchemaId()); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult all_rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(all_rows.data, EvolvedFields(), + R"([[0, 1, "old", null], [0, 2, "skip", null]])"); + + auto predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, "old", 3)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult filtered_rows, + ReadRows(table_path_, options, pool_, /*realtime_context=*/nullptr, + predicate, /*enable_predicate_filter=*/false)); + AssertResultEquals(filtered_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimeAppendNewWriterIndexesNewSchemaColumn) { + std::map options = BaseOptions(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); + + std::map options_v1 = options; + options_v1["file-index.bitmap.columns"] = "extra"; + options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options_v1)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[1, "new", 20], [2, "skip", 30]])", + /*partition=*/{}, /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> messages, + WriteWithNewWriter(table_path_, options_v1, std::move(new_schema_batch), + /*commit_identifier=*/1)); + ASSERT_EQ(1, messages.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(messages[0]); + ASSERT_TRUE(message); + std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(1, new_file->schema_id); + ASSERT_TRUE(new_file->embedded_index); + ASSERT_TRUE(new_file->extra_files.empty()); + ASSERT_OK_AND_ASSIGN( + std::vector> extra_indexes, + ReadEmbeddedIndexColumn(new_file, arrow::schema(EvolvedFields()), "extra", pool_)); + ASSERT_EQ(1, extra_indexes.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr extra_hit, + extra_indexes[0]->VisitEqual(Literal(20))); + ASSERT_OK_AND_ASSIGN(bool extra_remain, extra_hit->IsRemain()); + ASSERT_TRUE(extra_remain); + + ASSERT_OK(CommitMessages(table_path_, options_v1, messages, /*commit_identifier=*/1)); + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/2, /*field_name=*/"extra", FieldType::INT, Literal(20)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, + predicate, /*enable_predicate_filter=*/true)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimeAppendDataEvolutionWritesPartialNewColumnIndex) { + std::map options = DataEvolutionOptions(); + arrow::FieldVector fields = DataEvolutionFields(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(fields), /*primary_keys=*/{}, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeUnbucketedBatch(fields, R"([[1, "old", "base"]])", + /*partition=*/{})); + ASSERT_OK_AND_ASSIGN(std::vector> old_messages, + WriteWithNewWriter(table_path_, options, std::move(old_schema_batch), + /*commit_identifier=*/1, + /*write_schema=*/{"f0", "f1", "f2"})); + ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/1)); + + std::map options_v1 = DataEvolutionOptions(); + options_v1["file-index.bitmap.columns"] = "f2"; + options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; + ASSERT_OK( + EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, fields[0]), DataField(1, fields[1]), DataField(2, fields[2])}, + /*highest_field_id=*/2, options_v1)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr partial_batch, + MakeUnbucketedBatch({fields[2]}, R"([["updated"]])", + /*partition=*/{})); + ASSERT_OK_AND_ASSIGN(std::vector> new_messages, + WriteWithNewWriter(table_path_, options_v1, std::move(partial_batch), + /*commit_identifier=*/2, + /*write_schema=*/{"f2"})); + ASSERT_EQ(1, new_messages.size()); + std::shared_ptr new_message = + std::dynamic_pointer_cast(new_messages[0]); + ASSERT_TRUE(new_message); + std::shared_ptr new_file = new_message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(1, new_file->schema_id); + const std::optional> expected_write_cols = + std::vector{"f2"}; + ASSERT_EQ(expected_write_cols, new_file->write_cols); + ASSERT_TRUE(new_file->embedded_index); + ASSERT_TRUE(new_file->extra_files.empty()); + ASSERT_OK_AND_ASSIGN( + std::vector> f2_indexes, + ReadEmbeddedIndexColumn(new_file, arrow::schema({fields[2]}), "f2", pool_)); + ASSERT_EQ(1, f2_indexes.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr f2_hit, + f2_indexes[0]->VisitEqual(Literal(FieldType::STRING, "updated", 7))); + ASSERT_OK_AND_ASSIGN(bool f2_remain, f2_hit->IsRemain()); + ASSERT_TRUE(f2_remain); + + AssignFirstRowId(new_messages, /*first_row_id=*/0); + ASSERT_OK(CommitMessages(table_path_, options_v1, new_messages, /*commit_identifier=*/2)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult all_rows, ReadRows(table_path_, options_v1, pool_)); + AssertResultEquals(all_rows.data, fields, R"([[0, 1, "old", "updated"]])"); + + auto predicate = + PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2", FieldType::STRING, + Literal(FieldType::STRING, "updated", 7)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult filtered_rows, + ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, + predicate, /*enable_predicate_filter=*/false)); + AssertResultEquals(filtered_rows.data, fields, R"([[0, 1, "old", "updated"]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkOldAndNewSchemaFilesReadThroughLatestSchema) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> base_messages, + old_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/1)); + ASSERT_OK(CommitMessages(table_path_, options, base_messages, /*commit_identifier=*/1)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, + CreateWriter(table_path_, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> new_messages, + new_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_EQ(1, new_messages.size()); + std::shared_ptr new_message = + std::dynamic_pointer_cast(new_messages[0]); + ASSERT_TRUE(new_message); + ASSERT_EQ(1, new_message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_EQ(1, new_message->GetNewFilesIncrement().NewFiles()[0]->schema_id); + ASSERT_TRUE(new_message->GetNewFilesIncrement().NewFiles()[0]->extra_files.empty()); + + ASSERT_OK(CommitMessages(table_path_, options, new_messages, /*commit_identifier=*/2)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr stale_schema_batch, + MakeBatch(BaseFields(), R"([[3, "stale"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(stale_schema_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> stale_messages, + old_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/3)); + ASSERT_EQ(1, stale_messages.size()); + std::shared_ptr stale_message = + std::dynamic_pointer_cast(stale_messages[0]); + ASSERT_TRUE(stale_message); + ASSERT_EQ(1, stale_message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_EQ(0, stale_message->GetNewFilesIncrement().NewFiles()[0]->schema_id); + ASSERT_OK(CommitMessages(table_path_, options, stale_messages, /*commit_identifier=*/3)); + + ASSERT_OK_AND_ASSIGN(std::optional snapshot, + LatestSnapshot(table_path_, options, dir_->GetFileSystem())); + ASSERT_TRUE(snapshot); + ASSERT_EQ(1, snapshot->SchemaId()); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(rows.data, EvolvedFields(), + R"([[0, 1, "old", null], [0, 2, "new", 20], [0, 3, "stale", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkOldWriterCanOverwriteNewColumnWithNull) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[1, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> new_messages, + WriteWithNewWriter(table_path_, options, std::move(new_schema_batch), + /*commit_identifier=*/1)); + ASSERT_OK(CommitMessages(table_path_, options, new_messages, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> old_messages, + old_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/2)); + ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/2)); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkNewWriterIndexesNewSchemaColumn) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> old_messages, + WriteWithNewWriter(table_path_, options, std::move(old_schema_batch), + /*commit_identifier=*/1)); + ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/1)); + + std::map options_v1 = options; + options_v1["file-index.bitmap.columns"] = "extra"; + options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options_v1)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20], [3, "skip", 30]])", + /*partition=*/{}, /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> new_messages, + WriteWithNewWriter(table_path_, options_v1, std::move(new_schema_batch), + /*commit_identifier=*/2)); + ASSERT_EQ(1, new_messages.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(new_messages[0]); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(1, new_file->schema_id); + ASSERT_FALSE(new_file->embedded_index); + ASSERT_EQ(1, new_file->extra_files.size()); + ASSERT_TRUE(new_file->extra_files[0]); + std::string index_path = + PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); + ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); + ASSERT_TRUE(index_exists); + + ASSERT_OK(CommitMessages(table_path_, options_v1, new_messages, /*commit_identifier=*/2)); + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/2, /*field_name=*/"extra", FieldType::INT, Literal(20)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, + predicate, /*enable_predicate_filter=*/true)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 2, "new", 20]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkEmbeddedFileIndexFailsValueScan) { + std::map options = BaseOptions(); + options["file-index.bitmap.columns"] = "payload"; + options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(BaseFields(), R"([[1, "a"], [2, "b"]])", + /*partition=*/{}, /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> messages, + WriteWithNewWriter(table_path_, options, std::move(batch), + /*commit_identifier=*/1)); + ASSERT_EQ(1, messages.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(messages[0]); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_TRUE(message->GetNewFilesIncrement().NewFiles()[0]->embedded_index); + ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); + + std::shared_ptr predicate = + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, "a", 1)); + ASSERT_NOK_WITH_MSG( + ScanTable(table_path_, options, pool_, /*realtime_context=*/nullptr, predicate), + "do not support embedded index in DataFileMeta"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkSortedIndexConfigDoesNotWriteIndexOnDataWrite) { + std::map options = BaseOptions(); + options[Options::PK_BTREE_INDEX_COLUMNS] = "payload"; + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(BaseFields(), R"([[1, "a"], [2, "b"]])", + /*partition=*/{}, /*bucket=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector> messages, + WriteWithNewWriter(table_path_, options, std::move(batch), + /*commit_identifier=*/1)); + ASSERT_EQ(1, messages.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(messages[0]); + ASSERT_TRUE(message); + ASSERT_TRUE(message->GetNewFilesIncrement().NewIndexFiles().empty()); + ASSERT_TRUE(message->GetCompactIncrement().NewIndexFiles().empty()); + ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::optional snapshot, + LatestSnapshot(table_path_, options, dir_->GetFileSystem())); + ASSERT_TRUE(snapshot); + ASSERT_FALSE(snapshot->IndexManifest()); + + std::shared_ptr predicate = + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, "a", 1)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + ScanTable(table_path_, options, pool_, + /*realtime_context=*/nullptr, predicate)); + ASSERT_EQ(0, CountIndexedSplits(plan)); + std::vector> planned_files = DataFilesFromPlan(plan); + ASSERT_EQ(1, planned_files.size()); + ASSERT_EQ(0, planned_files[0]->schema_id); + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options, pool_, /*realtime_context=*/nullptr, + predicate, /*enable_predicate_filter=*/true)); + AssertResultEquals(rows.data, BaseFields(), R"([[0, 1, "a"]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendRejectsDataEvolutionAtWriterCreation) { + std::map create_options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, create_options)); + + std::map write_options = BaseOptions(); + write_options[Options::DATA_EVOLUTION_ENABLED] = "true"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_NOK_WITH_MSG(CreateWriter(table_path_, write_options, realtime_context), + "real-time append write does not support data evolution"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendScanRejectsDataEvolutionTable) { + std::map options = DataEvolutionOptions(); + arrow::FieldVector fields = DataEvolutionFields(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(fields), /*primary_keys=*/{}, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeUnbucketedBatch(fields, R"([[1, "old", "base"]])", + /*partition=*/{})); + ASSERT_OK_AND_ASSIGN(std::vector> messages, + WriteWithNewWriter(table_path_, options, std::move(batch), + /*commit_identifier=*/1, + /*write_schema=*/{"f0", "f1", "f2"})); + ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + + ASSERT_NOK_WITH_MSG(ScanTable(table_path_, options, pool_, realtime_context), + "real-time union read requires fixed bucket mode"); + + std::map fixed_bucket_options = options; + fixed_bucket_options[Options::BUCKET] = "1"; + ASSERT_NOK_WITH_MSG(ScanTable(table_path_, fixed_bucket_options, pool_, realtime_context), + "real-time union read does not support data evolution"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendReuseContextKeepsOldMemorySchemaAfterAlter) { + std::map options = BaseOptions(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + CommitRealtimeMessages(table_path_, options, base_progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reused_context_writer, + CreateWriter(table_path_, options, realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_NOK_WITH_MSG(reused_context_writer->Write(std::move(new_schema_batch)), + "ArrowArray struct has 3 children, expected 2"); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options, pool_, realtime_context)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendOldWriterAfterAlterCommitsOldSchemaFile) { + std::map options = BaseOptions(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, realtime_context)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); + ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, + ReadRows(table_path_, options, pool_, realtime_context)); + AssertResultEquals(memory_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(progress[0].commit_message); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_EQ(0, message->GetNewFilesIncrement().NewFiles()[0]->schema_id); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, CommitRealtimeMessages(table_path_, options, progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::optional snapshot, + LatestSnapshot(table_path_, options, dir_->GetFileSystem())); + ASSERT_TRUE(snapshot); + ASSERT_EQ(1, snapshot->SchemaId()); + ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendNewContextUsesNewSchemaAfterAlter) { + std::map options = BaseOptions(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr old_realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, old_realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + CommitRealtimeMessages(table_path_, options, base_progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + std::map options_v1 = options; + options_v1["file-index.bitmap.columns"] = "extra"; + options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, + CreateWriter(table_path_, options_v1, new_realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); + ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, + ReadRows(table_path_, options_v1, pool_, new_realtime_context)); + AssertResultEquals(memory_rows.data, EvolvedFields(), + R"([[0, 1, "old", null], [0, 2, "new", 20]])"); + + ASSERT_OK_AND_ASSIGN(std::vector new_progress, + new_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, new_progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(new_progress[0].commit_message); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(1, new_file->schema_id); + ASSERT_FALSE(new_file->embedded_index); + ASSERT_EQ(1, new_file->extra_files.size()); + ASSERT_TRUE(new_file->extra_files[0]); + std::string index_path = + PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); + ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); + ASSERT_TRUE(index_exists); + + ASSERT_OK_AND_ASSIGN(snapshot_id, CommitRealtimeMessages(table_path_, options_v1, new_progress, + /*commit_identifier=*/2)); + ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options_v1, pool_, new_realtime_context)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null], [0, 2, "new", 20]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkRejectsDataEvolutionAtWriterCreation) { + std::map create_options = BaseOptions(); + ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, + create_options)); + + std::map write_options = BaseOptions(); + write_options[Options::DATA_EVOLUTION_ENABLED] = "true"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_NOK_WITH_MSG(CreateWriter(table_path_, write_options, realtime_context), + "PK realtime v1 does not support data evolution"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkReuseContextKeepsOldMemorySchemaAfterAlter) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + CommitRealtimeMessages(table_path_, options, base_progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reused_context_writer, + CreateWriter(table_path_, options, realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_NOK_WITH_MSG(reused_context_writer->Write(std::move(new_schema_batch)), + "ArrowArray struct has 3 children, expected 2"); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options, pool_, realtime_context)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkOldWriterAfterAlterReadsNewColumnAsNull) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, realtime_context)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); + ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, + ReadRows(table_path_, options, pool_, realtime_context)); + AssertResultEquals(memory_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(progress[0].commit_message); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_EQ(0, message->GetNewFilesIncrement().NewFiles()[0]->schema_id); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, CommitRealtimeMessages(table_path_, options, progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, + RealtimePkOldWriterAfterAlterCannotCommitBehindNewContextOffset) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, realtime_context)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, + CreateWriter(table_path_, options, new_realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[1, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); + ASSERT_OK_AND_ASSIGN(std::vector new_progress, + new_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + CommitRealtimeMessages(table_path_, options, new_progress, + /*commit_identifier=*/1)); + ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); + ASSERT_OK_AND_ASSIGN(CollectedReadResult realtime_rows, + ReadRows(table_path_, options, pool_, realtime_context)); + AssertResultEquals(realtime_rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); + ASSERT_OK_AND_ASSIGN(std::vector old_progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_NOK_WITH_MSG(CommitRealtimeMessages(table_path_, options, old_progress, + /*commit_identifier=*/2), + "real-time commit offsets for bucket 0 are not contiguous"); + + ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); + AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); +} + +TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkNewContextUsesNewSchemaAfterAlter) { + std::map options = BaseOptions(); + ASSERT_OK( + CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr old_realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, + CreateWriter(table_path_, options, old_realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(old_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + CommitRealtimeMessages(table_path_, options, base_progress, + /*commit_identifier=*/1)); + ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), + {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), + DataField(2, EvolvedFields()[2])}, + /*highest_field_id=*/2, options)); + + std::map options_v1 = options; + options_v1["file-index.bitmap.columns"] = "extra"; + options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, + CreateWriter(table_path_, options_v1, new_realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, + MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, + /*bucket=*/0)); + ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); + ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, + ReadRows(table_path_, options_v1, pool_, new_realtime_context)); + AssertResultEquals(memory_rows.data, EvolvedFields(), + R"([[0, 1, "old", null], [0, 2, "new", 20]])"); + + ASSERT_OK_AND_ASSIGN(std::vector new_progress, + new_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, new_progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(new_progress[0].commit_message); + ASSERT_TRUE(message); + ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); + std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(1, new_file->schema_id); + ASSERT_FALSE(new_file->embedded_index); + ASSERT_EQ(1, new_file->extra_files.size()); + ASSERT_TRUE(new_file->extra_files[0]); + std::string index_path = + PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); + ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); + ASSERT_TRUE(index_exists); + + ASSERT_OK_AND_ASSIGN(snapshot_id, CommitRealtimeMessages(table_path_, options_v1, new_progress, + /*commit_identifier=*/2)); + ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, + ReadRows(table_path_, options_v1, pool_, new_realtime_context)); + AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null], [0, 2, "new", 20]])"); +} + +} // namespace +} // namespace paimon::test From 33583044237315a669ae727b1c7420b420616cbd Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Fri, 21 Aug 2026 14:41:09 +0800 Subject: [PATCH 12/24] fix(realtime): align PK projections by field ID --- .../realtime/primary_key_realtime_store.cpp | 36 +- .../primary_key_realtime_store_test.cpp | 50 +- test/inte/CMakeLists.txt | 7 - ...chema_evolution_write_verify_inte_test.cpp | 1110 ----------------- 4 files changed, 76 insertions(+), 1127 deletions(-) delete mode 100644 test/inte/schema_evolution_write_verify_inte_test.cpp diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 6565ed8d7..8f51c1b1e 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -98,6 +98,30 @@ uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { return result; } +int32_t FindPkQueryFieldIndex(const std::shared_ptr& write_schema, + const std::shared_ptr& read_field) { + Result read_id = NestedProjectionUtils::GetPaimonFieldId(read_field); + if (read_id.ok()) { + Result> write_field = + NestedProjectionUtils::FindFieldByPaimonId(arrow::struct_(write_schema->fields()), + read_id.value()); + if (write_field.ok()) { + return write_schema->GetFieldIndex(write_field.value()->name()); + } + } + + const int32_t name_index = write_schema->GetFieldIndex(read_field->name()); + if (name_index < 0) { + return -1; + } + Result write_id = + NestedProjectionUtils::GetPaimonFieldId(write_schema->field(name_index)); + if (read_id.ok() && write_id.ok() && read_id.value() != write_id.value()) { + return -1; + } + return name_index; +} + struct StoredBatch { std::shared_ptr data; std::vector row_kinds; @@ -423,17 +447,23 @@ class PrimaryKeyRealtimeStore::Impl { projection.push_back(KeyValueProjectionConsumer::kSequenceNumberProjection); continue; } - int32_t index = write_schema_->GetFieldIndex(field->name()); + int32_t index = FindPkQueryFieldIndex(write_schema_, field); if (index < 0) { Result field_id = NestedProjectionUtils::GetPaimonFieldId(field); if (!field_id.ok()) { return Status::Invalid( "PK real-time query field is missing from write schema: ", field->name()); } + std::string internal_name = + "__paimon_pk_realtime_null_" + std::to_string(field_id.value()); + while ( + NestedProjectionUtils::FindFieldByName(aligned_value_fields, internal_name)) { + internal_name.push_back('_'); + } index = static_cast(aligned_value_fields.size()); - aligned_value_fields.push_back(field); + aligned_value_fields.push_back(field->WithName(internal_name)); } else { - aligned_value_fields[index] = field; + aligned_value_fields[index] = write_schema_->field(index)->WithType(field->type()); } projection.push_back(index); } diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index ef293e54b..66901a6b1 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -392,6 +392,40 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjection) { "query field is missing from write schema: unknown"); } +TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjectionMatchesRenamedFieldsById) { + const std::shared_ptr id = + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))); + const std::shared_ptr value = + DataField::ConvertDataFieldToArrowField(DataField(1, arrow::field("value", arrow::utf8()))); + const std::shared_ptr write_schema = arrow::schema({id, value}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(write_schema, {"id"}, /*restore_max_sequence=*/4)); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[1, "kept"]])", {}, write_schema), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + const std::shared_ptr renamed_value = DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("renamed", arrow::utf8()))); + const std::shared_ptr renamed_id = DataField::ConvertDataFieldToArrowField( + DataField(0, arrow::field("renamed_id", arrow::int64()))); + const std::shared_ptr replaced = + DataField::ConvertDataFieldToArrowField(DataField(2, arrow::field("value", arrow::utf8()))); + const std::shared_ptr replaced_id = + DataField::ConvertDataFieldToArrowField(DataField(4, arrow::field("id", arrow::int64()))); + const std::shared_ptr added = + DataField::ConvertDataFieldToArrowField(DataField(3, arrow::field("added", arrow::utf8()))); + std::unique_ptr read_schema = + MakeReadSchema({renamed_value, renamed_id, replaced, replaced_id, added}); + RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + const std::shared_ptr result_type = + arrow::struct_({DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), + renamed_value, renamed_id, replaced, replaced_id, added}); + AssertReaderOutput(readers, result_type, R"([[0, "kept", 1, null, null, null]])"); +} + TEST_F(PrimaryKeyRealtimeStoreTest, TestNestedProjection) { const std::shared_ptr id = DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))); @@ -433,7 +467,10 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestCompositeKeyClipping) { {}, composite_schema), OffsetRange(20, 24)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields(composite_schema)); + const std::shared_ptr sequence = + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()); + std::unique_ptr read_schema = + MakeReadSchema({sequence, composite_schema->field(0), composite_schema->field(2)}); RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, @@ -445,13 +482,12 @@ TEST_F(PrimaryKeyRealtimeStoreTest, TestCompositeKeyClipping) { ASSERT_EQ("c", range->GetMinKey()->GetString(1).ToString()); ASSERT_EQ(2, range->GetMaxKey()->GetLong(0)); ASSERT_EQ("b", range->GetMaxKey()->GetString(1).ToString()); - std::shared_ptr query_type = arrow::struct_( - {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()), - composite_schema->field(0), composite_schema->field(1), composite_schema->field(2)}); + std::shared_ptr query_type = + arrow::struct_({DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), + sequence, composite_schema->field(0), composite_schema->field(2)}); AssertReaderOutput(readers, query_type, - R"([[0, 7, 1, "c", "one-c"], [0, 8, 2, "a", "two-a"], - [0, 6, 2, "b", "two-b"]])"); + R"([[0, 7, 1, "one-c"], [0, 8, 2, "two-a"], + [0, 6, 2, "two-b"]])"); } } // namespace paimon::test diff --git a/test/inte/CMakeLists.txt b/test/inte/CMakeLists.txt index f1b3f8ce6..75147ce60 100644 --- a/test/inte/CMakeLists.txt +++ b/test/inte/CMakeLists.txt @@ -43,13 +43,6 @@ if(PAIMON_BUILD_TESTS) test_utils_static ${GTEST_LINK_TOOLCHAIN}) - add_paimon_test(schema_evolution_write_verify_inte_test - STATIC_LINK_LIBS - paimon_shared - ${TEST_STATIC_LINK_LIBS} - test_utils_static - ${GTEST_LINK_TOOLCHAIN}) - add_paimon_test(global_index_test STATIC_LINK_LIBS paimon_shared diff --git a/test/inte/schema_evolution_write_verify_inte_test.cpp b/test/inte/schema_evolution_write_verify_inte_test.cpp deleted file mode 100644 index dcadbd9e1..000000000 --- a/test/inte/schema_evolution_write_verify_inte_test.cpp +++ /dev/null @@ -1,1110 +0,0 @@ -/* - * 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 -#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/common/utils/scope_guard.h" -#include "paimon/core/core_options.h" -#include "paimon/core/global_index/indexed_split_impl.h" -#include "paimon/core/io/data_file_meta.h" -#include "paimon/core/schema/schema_manager.h" -#include "paimon/core/snapshot.h" -#include "paimon/core/table/sink/commit_message_impl.h" -#include "paimon/core/table/source/data_split_impl.h" -#include "paimon/defs.h" -#include "paimon/file_index/file_index_format.h" -#include "paimon/file_index/file_index_reader.h" -#include "paimon/file_index/file_index_result.h" -#include "paimon/file_store_commit.h" -#include "paimon/file_store_write.h" -#include "paimon/fs/file_system.h" -#include "paimon/io/byte_array_input_stream.h" -#include "paimon/predicate/literal.h" -#include "paimon/predicate/predicate_builder.h" -#include "paimon/read_context.h" -#include "paimon/reader/batch_reader.h" -#include "paimon/realtime/realtime_context.h" -#include "paimon/record_batch.h" -#include "paimon/scan_context.h" -#include "paimon/table/source/plan.h" -#include "paimon/table/source/startup_mode.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/test_helper.h" -#include "paimon/testing/utils/testharness.h" -#include "paimon/write_context.h" - -namespace paimon::test { -namespace { - -std::map BaseOptions() { - return {{Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, - {Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "1"}, - {Options::BUCKET_KEY, "id"}, {Options::TARGET_FILE_SIZE, "1MB"}}; -} - -std::map DataEvolutionOptions() { - return {{Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, - {Options::FILE_SYSTEM, "local"}, {Options::TARGET_FILE_SIZE, "1MB"}, - {Options::ROW_TRACKING_ENABLED, "true"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}; -} - -arrow::FieldVector BaseFields() { - return {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8())}; -} - -arrow::FieldVector EvolvedFields() { - return {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), - arrow::field("extra", arrow::int32())}; -} - -arrow::FieldVector DataEvolutionFields() { - return {arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::utf8()), - arrow::field("f2", arrow::utf8())}; -} - -Result> MakeBatch( - const arrow::FieldVector& fields, const std::string& json, - const std::map& partition, int32_t bucket, - const std::vector& row_kinds = {}) { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr array, - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), json)); - ArrowArray c_array; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - RecordBatchBuilder builder(&c_array); - return builder.SetPartition(partition).SetBucket(bucket).SetRowKinds(row_kinds).Finish(); -} - -Result> MakeUnbucketedBatch( - const arrow::FieldVector& fields, const std::string& json, - const std::map& partition, - const std::vector& row_kinds = {}) { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr array, - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), json)); - ArrowArray c_array; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); - RecordBatchBuilder builder(&c_array); - return builder.SetPartition(partition).SetRowKinds(row_kinds).Finish(); -} - -Result> CreateWriter( - const std::string& table_path, const std::map& options, - const std::shared_ptr& realtime_context = nullptr, - const std::vector& write_schema = {}) { - WriteContextBuilder builder(table_path, "schema_evolution_verify"); - builder.SetOptions(options).WithStreamingMode(true); - if (realtime_context) { - builder.WithRealtimeContext(realtime_context); - } - if (!write_schema.empty()) { - builder.WithWriteSchema(write_schema); - } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); - return FileStoreWrite::Create(std::move(context)); -} - -Result>> WriteWithNewWriter( - const std::string& table_path, const std::map& options, - std::unique_ptr batch, int64_t commit_identifier, - const std::vector& write_schema = {}) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, - CreateWriter(table_path, options, nullptr, write_schema)); - 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> CreateCommit( - const std::string& table_path, const std::map& options) { - CommitContextBuilder builder(table_path, "schema_evolution_verify"); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, - builder.SetOptions(options).IgnoreEmptyCommit(false).Finish()); - return FileStoreCommit::Create(std::move(context)); -} - -Status CommitMessages(const std::string& table_path, - const std::map& options, - const std::vector>& messages, - int64_t commit_identifier) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, - CreateCommit(table_path, options)); - return commit->Commit(messages, commit_identifier); -} - -Result CommitRealtimeMessages(const std::string& table_path, - const std::map& options, - const std::vector& messages, - int64_t commit_identifier) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, - CreateCommit(table_path, options)); - return commit->CommitWithProgress(messages, commit_identifier, /*watermark=*/std::nullopt); -} - -Result> LatestSnapshot(const std::string& table_path, - const std::map& options, - const std::shared_ptr& file_system) { - PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options, file_system)); - SnapshotManager snapshot_manager(core_options.GetFileSystem(), table_path); - return snapshot_manager.LatestSnapshot(); -} - -Result> ScanTable( - const std::string& table_path, const std::map& options, - const std::shared_ptr& pool, - const std::shared_ptr& realtime_context = nullptr, - const std::shared_ptr& predicate = nullptr) { - ScanContextBuilder scan_builder(table_path); - scan_builder.SetOptions(options) - .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString()) - .SetPredicate(predicate) - .WithMemoryPool(pool); - if (realtime_context) { - scan_builder.WithRealtimeContext(realtime_context); - } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, scan_builder.Finish()); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_scan, - TableScan::Create(std::move(scan_context))); - return table_scan->CreatePlan(); -} - -std::vector> DataFilesFromPlan(const std::shared_ptr& plan) { - std::vector> files; - for (const std::shared_ptr& split : plan->Splits()) { - std::shared_ptr data_split = split; - if (std::shared_ptr indexed_split = - std::dynamic_pointer_cast(split)) { - data_split = indexed_split->GetDataSplit(); - } - std::shared_ptr split_impl = - std::dynamic_pointer_cast(data_split); - if (!split_impl) { - continue; - } - const std::vector>& split_files = split_impl->DataFiles(); - files.insert(files.end(), split_files.begin(), split_files.end()); - } - return files; -} - -size_t CountIndexedSplits(const std::shared_ptr& plan) { - size_t count = 0; - for (const std::shared_ptr& split : plan->Splits()) { - if (std::dynamic_pointer_cast(split)) { - count++; - } - } - return count; -} - -Status EvolveSchema(const std::string& table_path, const std::shared_ptr& file_system, - const std::vector& fields, int32_t highest_field_id, - const std::map& options) { - return TestHelper::WriteNextSchema(file_system, table_path, fields, highest_field_id, options); -} - -void AssignFirstRowId(const std::vector>& messages, - int64_t first_row_id) { - for (const std::shared_ptr& commit_message : messages) { - std::shared_ptr message = - std::dynamic_pointer_cast(commit_message); - ASSERT_TRUE(message); - for (const std::shared_ptr& file : - message->GetNewFilesIncrement().NewFiles()) { - file->AssignFirstRowId(first_row_id); - } - } -} - -struct CollectedReadResult { - std::unique_ptr table_read; - std::unique_ptr reader; - std::shared_ptr data; -}; - -Result ReadRows( - const std::string& table_path, const std::map& options, - const std::shared_ptr& pool, - const std::shared_ptr& realtime_context = nullptr, - const std::shared_ptr& predicate = nullptr, bool enable_predicate_filter = true) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, - ScanTable(table_path, options, pool, realtime_context, predicate)); - - ReadContextBuilder read_builder(table_path); - read_builder.SetOptions(options) - .SetPredicate(predicate) - .EnablePredicateFilter(enable_predicate_filter) - .WithMemoryPool(pool); - if (realtime_context) { - read_builder.WithRealtimeContext(realtime_context); - } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, - TableRead::Create(std::move(read_context))); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr batch_reader, - table_read->CreateReader(plan->Splits())); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr rows, - ReadResultCollector::CollectResult(batch_reader.get())); - return CollectedReadResult{std::move(table_read), std::move(batch_reader), std::move(rows)}; -} - -void AssertResultEquals(const std::shared_ptr& actual, - const arrow::FieldVector& fields, const std::string& expected_json) { - arrow::FieldVector fields_with_row_kind = fields; - fields_with_row_kind.insert(fields_with_row_kind.begin(), - arrow::field("_VALUE_KIND", arrow::int8())); - std::shared_ptr expected_array = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_with_row_kind), - expected_json) - .ValueOrDie(); - auto expected = std::make_shared(expected_array); - ASSERT_TRUE(expected->Equals(actual, arrow::EqualOptions::Defaults().diff_sink(&std::cout))) - << actual->ToString(); -} - -Status CreateTable(const std::string& warehouse, const std::shared_ptr& schema, - const std::vector& primary_keys, - const std::map& options) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr catalog, Catalog::Create(warehouse, options)); - PAIMON_RETURN_NOT_OK(catalog->CreateDatabase("foo", options, /*ignore_if_exists=*/false)); - ArrowSchema c_schema; - ArrowSchemaMarkReleased(&c_schema); - ScopeGuard guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); - return catalog->CreateTable(Identifier("foo", "bar"), &c_schema, - /*partition_keys=*/{}, primary_keys, options, - /*ignore_if_exists=*/false); -} - -Result> CreateFileIndexReader( - const std::shared_ptr& data_file, const std::shared_ptr& pool) { - if (data_file->embedded_index == nullptr) { - return Status::Invalid("data file does not contain an embedded file index"); - } - auto input = std::make_shared(data_file->embedded_index->data(), - data_file->embedded_index->size()); - return FileIndexFormat::CreateReader(input, pool); -} - -Result>> ReadEmbeddedIndexColumn( - const std::shared_ptr& data_file, const std::shared_ptr& schema, - const std::string& column, const std::shared_ptr& pool) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, - CreateFileIndexReader(data_file, pool)); - auto c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, c_schema.get())); - return reader->ReadColumnIndex(column, c_schema.get()); -} - -class SchemaEvolutionWriteVerifyTest : public ::testing::Test { - protected: - void SetUp() override { - pool_ = GetDefaultPool(); - dir_ = UniqueTestDirectory::Create("local"); - ASSERT_TRUE(dir_); - table_path_ = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); - } - - void TearDown() override { - dir_.reset(); - } - - std::shared_ptr pool_; - std::unique_ptr dir_; - std::string table_path_; -}; - -TEST_F(SchemaEvolutionWriteVerifyTest, - NonRealtimeAppendOldWriterCommitsOldSchemaFileIntoNewSchemaSnapshot) { - std::map options = BaseOptions(); - options["file-index.bitmap.columns"] = "payload"; - options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"], [2, "skip"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); - ASSERT_OK_AND_ASSIGN(std::vector> messages, - old_writer->PrepareCommit(/*wait_compaction=*/false, - /*commit_identifier=*/1)); - ASSERT_EQ(1, messages.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(messages[0]); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - std::shared_ptr old_file = message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(0, old_file->schema_id); - ASSERT_TRUE(old_file->embedded_index); - ASSERT_TRUE(old_file->extra_files.empty()); - ASSERT_OK_AND_ASSIGN( - std::vector> payload_indexes, - ReadEmbeddedIndexColumn(old_file, arrow::schema(BaseFields()), "payload", pool_)); - ASSERT_EQ(1, payload_indexes.size()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr payload_hit, - payload_indexes[0]->VisitEqual(Literal(FieldType::STRING, "old", 3))); - ASSERT_OK_AND_ASSIGN(bool payload_remain, payload_hit->IsRemain()); - ASSERT_TRUE(payload_remain); - - ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(std::optional snapshot, - LatestSnapshot(table_path_, options, dir_->GetFileSystem())); - ASSERT_TRUE(snapshot); - ASSERT_EQ(1, snapshot->SchemaId()); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult all_rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(all_rows.data, EvolvedFields(), - R"([[0, 1, "old", null], [0, 2, "skip", null]])"); - - auto predicate = PredicateBuilder::Equal( - /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, - Literal(FieldType::STRING, "old", 3)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult filtered_rows, - ReadRows(table_path_, options, pool_, /*realtime_context=*/nullptr, - predicate, /*enable_predicate_filter=*/false)); - AssertResultEquals(filtered_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimeAppendNewWriterIndexesNewSchemaColumn) { - std::map options = BaseOptions(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); - - std::map options_v1 = options; - options_v1["file-index.bitmap.columns"] = "extra"; - options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options_v1)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[1, "new", 20], [2, "skip", 30]])", - /*partition=*/{}, /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> messages, - WriteWithNewWriter(table_path_, options_v1, std::move(new_schema_batch), - /*commit_identifier=*/1)); - ASSERT_EQ(1, messages.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(messages[0]); - ASSERT_TRUE(message); - std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(1, new_file->schema_id); - ASSERT_TRUE(new_file->embedded_index); - ASSERT_TRUE(new_file->extra_files.empty()); - ASSERT_OK_AND_ASSIGN( - std::vector> extra_indexes, - ReadEmbeddedIndexColumn(new_file, arrow::schema(EvolvedFields()), "extra", pool_)); - ASSERT_EQ(1, extra_indexes.size()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr extra_hit, - extra_indexes[0]->VisitEqual(Literal(20))); - ASSERT_OK_AND_ASSIGN(bool extra_remain, extra_hit->IsRemain()); - ASSERT_TRUE(extra_remain); - - ASSERT_OK(CommitMessages(table_path_, options_v1, messages, /*commit_identifier=*/1)); - std::shared_ptr predicate = PredicateBuilder::Equal( - /*field_index=*/2, /*field_name=*/"extra", FieldType::INT, Literal(20)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, - predicate, /*enable_predicate_filter=*/true)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimeAppendDataEvolutionWritesPartialNewColumnIndex) { - std::map options = DataEvolutionOptions(); - arrow::FieldVector fields = DataEvolutionFields(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(fields), /*primary_keys=*/{}, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeUnbucketedBatch(fields, R"([[1, "old", "base"]])", - /*partition=*/{})); - ASSERT_OK_AND_ASSIGN(std::vector> old_messages, - WriteWithNewWriter(table_path_, options, std::move(old_schema_batch), - /*commit_identifier=*/1, - /*write_schema=*/{"f0", "f1", "f2"})); - ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/1)); - - std::map options_v1 = DataEvolutionOptions(); - options_v1["file-index.bitmap.columns"] = "f2"; - options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; - ASSERT_OK( - EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, fields[0]), DataField(1, fields[1]), DataField(2, fields[2])}, - /*highest_field_id=*/2, options_v1)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr partial_batch, - MakeUnbucketedBatch({fields[2]}, R"([["updated"]])", - /*partition=*/{})); - ASSERT_OK_AND_ASSIGN(std::vector> new_messages, - WriteWithNewWriter(table_path_, options_v1, std::move(partial_batch), - /*commit_identifier=*/2, - /*write_schema=*/{"f2"})); - ASSERT_EQ(1, new_messages.size()); - std::shared_ptr new_message = - std::dynamic_pointer_cast(new_messages[0]); - ASSERT_TRUE(new_message); - std::shared_ptr new_file = new_message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(1, new_file->schema_id); - const std::optional> expected_write_cols = - std::vector{"f2"}; - ASSERT_EQ(expected_write_cols, new_file->write_cols); - ASSERT_TRUE(new_file->embedded_index); - ASSERT_TRUE(new_file->extra_files.empty()); - ASSERT_OK_AND_ASSIGN( - std::vector> f2_indexes, - ReadEmbeddedIndexColumn(new_file, arrow::schema({fields[2]}), "f2", pool_)); - ASSERT_EQ(1, f2_indexes.size()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr f2_hit, - f2_indexes[0]->VisitEqual(Literal(FieldType::STRING, "updated", 7))); - ASSERT_OK_AND_ASSIGN(bool f2_remain, f2_hit->IsRemain()); - ASSERT_TRUE(f2_remain); - - AssignFirstRowId(new_messages, /*first_row_id=*/0); - ASSERT_OK(CommitMessages(table_path_, options_v1, new_messages, /*commit_identifier=*/2)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult all_rows, ReadRows(table_path_, options_v1, pool_)); - AssertResultEquals(all_rows.data, fields, R"([[0, 1, "old", "updated"]])"); - - auto predicate = - PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"f2", FieldType::STRING, - Literal(FieldType::STRING, "updated", 7)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult filtered_rows, - ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, - predicate, /*enable_predicate_filter=*/false)); - AssertResultEquals(filtered_rows.data, fields, R"([[0, 1, "old", "updated"]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkOldAndNewSchemaFilesReadThroughLatestSchema) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(base_batch))); - ASSERT_OK_AND_ASSIGN(std::vector> base_messages, - old_writer->PrepareCommit(/*wait_compaction=*/false, - /*commit_identifier=*/1)); - ASSERT_OK(CommitMessages(table_path_, options, base_messages, /*commit_identifier=*/1)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, - CreateWriter(table_path_, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); - ASSERT_OK_AND_ASSIGN(std::vector> new_messages, - new_writer->PrepareCommit(/*wait_compaction=*/false, - /*commit_identifier=*/2)); - ASSERT_EQ(1, new_messages.size()); - std::shared_ptr new_message = - std::dynamic_pointer_cast(new_messages[0]); - ASSERT_TRUE(new_message); - ASSERT_EQ(1, new_message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_EQ(1, new_message->GetNewFilesIncrement().NewFiles()[0]->schema_id); - ASSERT_TRUE(new_message->GetNewFilesIncrement().NewFiles()[0]->extra_files.empty()); - - ASSERT_OK(CommitMessages(table_path_, options, new_messages, /*commit_identifier=*/2)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr stale_schema_batch, - MakeBatch(BaseFields(), R"([[3, "stale"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(stale_schema_batch))); - ASSERT_OK_AND_ASSIGN(std::vector> stale_messages, - old_writer->PrepareCommit(/*wait_compaction=*/false, - /*commit_identifier=*/3)); - ASSERT_EQ(1, stale_messages.size()); - std::shared_ptr stale_message = - std::dynamic_pointer_cast(stale_messages[0]); - ASSERT_TRUE(stale_message); - ASSERT_EQ(1, stale_message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_EQ(0, stale_message->GetNewFilesIncrement().NewFiles()[0]->schema_id); - ASSERT_OK(CommitMessages(table_path_, options, stale_messages, /*commit_identifier=*/3)); - - ASSERT_OK_AND_ASSIGN(std::optional snapshot, - LatestSnapshot(table_path_, options, dir_->GetFileSystem())); - ASSERT_TRUE(snapshot); - ASSERT_EQ(1, snapshot->SchemaId()); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(rows.data, EvolvedFields(), - R"([[0, 1, "old", null], [0, 2, "new", 20], [0, 3, "stale", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkOldWriterCanOverwriteNewColumnWithNull) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[1, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> new_messages, - WriteWithNewWriter(table_path_, options, std::move(new_schema_batch), - /*commit_identifier=*/1)); - ASSERT_OK(CommitMessages(table_path_, options, new_messages, /*commit_identifier=*/1)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); - ASSERT_OK_AND_ASSIGN(std::vector> old_messages, - old_writer->PrepareCommit(/*wait_compaction=*/false, - /*commit_identifier=*/2)); - ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/2)); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkNewWriterIndexesNewSchemaColumn) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> old_messages, - WriteWithNewWriter(table_path_, options, std::move(old_schema_batch), - /*commit_identifier=*/1)); - ASSERT_OK(CommitMessages(table_path_, options, old_messages, /*commit_identifier=*/1)); - - std::map options_v1 = options; - options_v1["file-index.bitmap.columns"] = "extra"; - options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options_v1)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20], [3, "skip", 30]])", - /*partition=*/{}, /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> new_messages, - WriteWithNewWriter(table_path_, options_v1, std::move(new_schema_batch), - /*commit_identifier=*/2)); - ASSERT_EQ(1, new_messages.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(new_messages[0]); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(1, new_file->schema_id); - ASSERT_FALSE(new_file->embedded_index); - ASSERT_EQ(1, new_file->extra_files.size()); - ASSERT_TRUE(new_file->extra_files[0]); - std::string index_path = - PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); - ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); - ASSERT_TRUE(index_exists); - - ASSERT_OK(CommitMessages(table_path_, options_v1, new_messages, /*commit_identifier=*/2)); - std::shared_ptr predicate = PredicateBuilder::Equal( - /*field_index=*/2, /*field_name=*/"extra", FieldType::INT, Literal(20)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options_v1, pool_, /*realtime_context=*/nullptr, - predicate, /*enable_predicate_filter=*/true)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 2, "new", 20]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkEmbeddedFileIndexFailsValueScan) { - std::map options = BaseOptions(); - options["file-index.bitmap.columns"] = "payload"; - options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch(BaseFields(), R"([[1, "a"], [2, "b"]])", - /*partition=*/{}, /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> messages, - WriteWithNewWriter(table_path_, options, std::move(batch), - /*commit_identifier=*/1)); - ASSERT_EQ(1, messages.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(messages[0]); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_TRUE(message->GetNewFilesIncrement().NewFiles()[0]->embedded_index); - ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); - - std::shared_ptr predicate = - PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, - Literal(FieldType::STRING, "a", 1)); - ASSERT_NOK_WITH_MSG( - ScanTable(table_path_, options, pool_, /*realtime_context=*/nullptr, predicate), - "do not support embedded index in DataFileMeta"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, NonRealtimePkSortedIndexConfigDoesNotWriteIndexOnDataWrite) { - std::map options = BaseOptions(); - options[Options::PK_BTREE_INDEX_COLUMNS] = "payload"; - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch(BaseFields(), R"([[1, "a"], [2, "b"]])", - /*partition=*/{}, /*bucket=*/0)); - ASSERT_OK_AND_ASSIGN(std::vector> messages, - WriteWithNewWriter(table_path_, options, std::move(batch), - /*commit_identifier=*/1)); - ASSERT_EQ(1, messages.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(messages[0]); - ASSERT_TRUE(message); - ASSERT_TRUE(message->GetNewFilesIncrement().NewIndexFiles().empty()); - ASSERT_TRUE(message->GetCompactIncrement().NewIndexFiles().empty()); - ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(std::optional snapshot, - LatestSnapshot(table_path_, options, dir_->GetFileSystem())); - ASSERT_TRUE(snapshot); - ASSERT_FALSE(snapshot->IndexManifest()); - - std::shared_ptr predicate = - PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, - Literal(FieldType::STRING, "a", 1)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, - ScanTable(table_path_, options, pool_, - /*realtime_context=*/nullptr, predicate)); - ASSERT_EQ(0, CountIndexedSplits(plan)); - std::vector> planned_files = DataFilesFromPlan(plan); - ASSERT_EQ(1, planned_files.size()); - ASSERT_EQ(0, planned_files[0]->schema_id); - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options, pool_, /*realtime_context=*/nullptr, - predicate, /*enable_predicate_filter=*/true)); - AssertResultEquals(rows.data, BaseFields(), R"([[0, 1, "a"]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendRejectsDataEvolutionAtWriterCreation) { - std::map create_options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, create_options)); - - std::map write_options = BaseOptions(); - write_options[Options::DATA_EVOLUTION_ENABLED] = "true"; - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_NOK_WITH_MSG(CreateWriter(table_path_, write_options, realtime_context), - "real-time append write does not support data evolution"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendScanRejectsDataEvolutionTable) { - std::map options = DataEvolutionOptions(); - arrow::FieldVector fields = DataEvolutionFields(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(fields), /*primary_keys=*/{}, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeUnbucketedBatch(fields, R"([[1, "old", "base"]])", - /*partition=*/{})); - ASSERT_OK_AND_ASSIGN(std::vector> messages, - WriteWithNewWriter(table_path_, options, std::move(batch), - /*commit_identifier=*/1, - /*write_schema=*/{"f0", "f1", "f2"})); - ASSERT_OK(CommitMessages(table_path_, options, messages, /*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - - ASSERT_NOK_WITH_MSG(ScanTable(table_path_, options, pool_, realtime_context), - "real-time union read requires fixed bucket mode"); - - std::map fixed_bucket_options = options; - fixed_bucket_options[Options::BUCKET] = "1"; - ASSERT_NOK_WITH_MSG(ScanTable(table_path_, fixed_bucket_options, pool_, realtime_context), - "real-time union read does not support data evolution"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendReuseContextKeepsOldMemorySchemaAfterAlter) { - std::map options = BaseOptions(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(base_batch))); - ASSERT_OK_AND_ASSIGN(std::vector base_progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, - CommitRealtimeMessages(table_path_, options, base_progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr reused_context_writer, - CreateWriter(table_path_, options, realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_NOK_WITH_MSG(reused_context_writer->Write(std::move(new_schema_batch)), - "ArrowArray struct has 3 children, expected 2"); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options, pool_, realtime_context)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendOldWriterAfterAlterCommitsOldSchemaFile) { - std::map options = BaseOptions(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, realtime_context)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, - ReadRows(table_path_, options, pool_, realtime_context)); - AssertResultEquals(memory_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); - - ASSERT_OK_AND_ASSIGN(std::vector progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_EQ(1, progress.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(progress[0].commit_message); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_EQ(0, message->GetNewFilesIncrement().NewFiles()[0]->schema_id); - - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, CommitRealtimeMessages(table_path_, options, progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - ASSERT_OK_AND_ASSIGN(std::optional snapshot, - LatestSnapshot(table_path_, options, dir_->GetFileSystem())); - ASSERT_TRUE(snapshot); - ASSERT_EQ(1, snapshot->SchemaId()); - ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimeAppendNewContextUsesNewSchemaAfterAlter) { - std::map options = BaseOptions(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr old_realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, old_realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(base_batch))); - ASSERT_OK_AND_ASSIGN(std::vector base_progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, - CommitRealtimeMessages(table_path_, options, base_progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - std::map options_v1 = options; - options_v1["file-index.bitmap.columns"] = "extra"; - options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; - ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, - CreateWriter(table_path_, options_v1, new_realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, - ReadRows(table_path_, options_v1, pool_, new_realtime_context)); - AssertResultEquals(memory_rows.data, EvolvedFields(), - R"([[0, 1, "old", null], [0, 2, "new", 20]])"); - - ASSERT_OK_AND_ASSIGN(std::vector new_progress, - new_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); - ASSERT_EQ(1, new_progress.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(new_progress[0].commit_message); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(1, new_file->schema_id); - ASSERT_FALSE(new_file->embedded_index); - ASSERT_EQ(1, new_file->extra_files.size()); - ASSERT_TRUE(new_file->extra_files[0]); - std::string index_path = - PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); - ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); - ASSERT_TRUE(index_exists); - - ASSERT_OK_AND_ASSIGN(snapshot_id, CommitRealtimeMessages(table_path_, options_v1, new_progress, - /*commit_identifier=*/2)); - ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options_v1, pool_, new_realtime_context)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null], [0, 2, "new", 20]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkRejectsDataEvolutionAtWriterCreation) { - std::map create_options = BaseOptions(); - ASSERT_OK(CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, - create_options)); - - std::map write_options = BaseOptions(); - write_options[Options::DATA_EVOLUTION_ENABLED] = "true"; - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_NOK_WITH_MSG(CreateWriter(table_path_, write_options, realtime_context), - "PK realtime v1 does not support data evolution"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkReuseContextKeepsOldMemorySchemaAfterAlter) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(base_batch))); - ASSERT_OK_AND_ASSIGN(std::vector base_progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, - CommitRealtimeMessages(table_path_, options, base_progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr reused_context_writer, - CreateWriter(table_path_, options, realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_NOK_WITH_MSG(reused_context_writer->Write(std::move(new_schema_batch)), - "ArrowArray struct has 3 children, expected 2"); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options, pool_, realtime_context)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkOldWriterAfterAlterReadsNewColumnAsNull) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, realtime_context)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, - ReadRows(table_path_, options, pool_, realtime_context)); - AssertResultEquals(memory_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); - - ASSERT_OK_AND_ASSIGN(std::vector progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_EQ(1, progress.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(progress[0].commit_message); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - ASSERT_EQ(0, message->GetNewFilesIncrement().NewFiles()[0]->schema_id); - - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, CommitRealtimeMessages(table_path_, options, progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "old", null]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, - RealtimePkOldWriterAfterAlterCannotCommitBehindNewContextOffset) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, realtime_context)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, - CreateWriter(table_path_, options, new_realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[1, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); - ASSERT_OK_AND_ASSIGN(std::vector new_progress, - new_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, - CommitRealtimeMessages(table_path_, options, new_progress, - /*commit_identifier=*/1)); - ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_schema_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(old_schema_batch))); - ASSERT_OK_AND_ASSIGN(CollectedReadResult realtime_rows, - ReadRows(table_path_, options, pool_, realtime_context)); - AssertResultEquals(realtime_rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); - ASSERT_OK_AND_ASSIGN(std::vector old_progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); - ASSERT_NOK_WITH_MSG(CommitRealtimeMessages(table_path_, options, old_progress, - /*commit_identifier=*/2), - "real-time commit offsets for bucket 0 are not contiguous"); - - ASSERT_OK_AND_ASSIGN(CollectedReadResult disk_rows, ReadRows(table_path_, options, pool_)); - AssertResultEquals(disk_rows.data, EvolvedFields(), R"([[0, 1, "new", 20]])"); -} - -TEST_F(SchemaEvolutionWriteVerifyTest, RealtimePkNewContextUsesNewSchemaAfterAlter) { - std::map options = BaseOptions(); - ASSERT_OK( - CreateTable(dir_->Str(), arrow::schema(BaseFields()), /*primary_keys=*/{"id"}, options)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr old_realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr old_writer, - CreateWriter(table_path_, options, old_realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, - MakeBatch(BaseFields(), R"([[1, "old"]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(old_writer->Write(std::move(base_batch))); - ASSERT_OK_AND_ASSIGN(std::vector base_progress, - old_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, - CommitRealtimeMessages(table_path_, options, base_progress, - /*commit_identifier=*/1)); - ASSERT_OK(old_writer->RefreshCommittedSnapshot(snapshot_id)); - - ASSERT_OK(EvolveSchema(table_path_, dir_->GetFileSystem(), - {DataField(0, BaseFields()[0]), DataField(1, BaseFields()[1]), - DataField(2, EvolvedFields()[2])}, - /*highest_field_id=*/2, options)); - - std::map options_v1 = options; - options_v1["file-index.bitmap.columns"] = "extra"; - options_v1[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1B"; - ASSERT_OK_AND_ASSIGN(std::shared_ptr new_realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_writer, - CreateWriter(table_path_, options_v1, new_realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr new_schema_batch, - MakeBatch(EvolvedFields(), R"([[2, "new", 20]])", /*partition=*/{}, - /*bucket=*/0)); - ASSERT_OK(new_writer->Write(std::move(new_schema_batch))); - ASSERT_OK_AND_ASSIGN(CollectedReadResult memory_rows, - ReadRows(table_path_, options_v1, pool_, new_realtime_context)); - AssertResultEquals(memory_rows.data, EvolvedFields(), - R"([[0, 1, "old", null], [0, 2, "new", 20]])"); - - ASSERT_OK_AND_ASSIGN(std::vector new_progress, - new_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); - ASSERT_EQ(1, new_progress.size()); - std::shared_ptr message = - std::dynamic_pointer_cast(new_progress[0].commit_message); - ASSERT_TRUE(message); - ASSERT_EQ(1, message->GetNewFilesIncrement().NewFiles().size()); - std::shared_ptr new_file = message->GetNewFilesIncrement().NewFiles()[0]; - ASSERT_EQ(1, new_file->schema_id); - ASSERT_FALSE(new_file->embedded_index); - ASSERT_EQ(1, new_file->extra_files.size()); - ASSERT_TRUE(new_file->extra_files[0]); - std::string index_path = - PathUtil::JoinPath(table_path_, "bucket-0/" + new_file->extra_files[0].value()); - ASSERT_OK_AND_ASSIGN(bool index_exists, dir_->GetFileSystem()->Exists(index_path)); - ASSERT_TRUE(index_exists); - - ASSERT_OK_AND_ASSIGN(snapshot_id, CommitRealtimeMessages(table_path_, options_v1, new_progress, - /*commit_identifier=*/2)); - ASSERT_OK(new_writer->RefreshCommittedSnapshot(snapshot_id)); - ASSERT_OK_AND_ASSIGN(CollectedReadResult rows, - ReadRows(table_path_, options_v1, pool_, new_realtime_context)); - AssertResultEquals(rows.data, EvolvedFields(), R"([[0, 1, "old", null], [0, 2, "new", 20]])"); -} - -} // namespace -} // namespace paimon::test From 8d961530efa9ce41b8885d9ed17160ac5dcdcaa2 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:21:45 +0800 Subject: [PATCH 13/24] refactor(mergetree): accept sorted key-value readers --- .../core/mergetree/merge_tree_writer.cpp | 95 ++++--- src/paimon/core/mergetree/merge_tree_writer.h | 3 + .../core/mergetree/merge_tree_writer_test.cpp | 236 ++++++++++++++++++ 3 files changed, 293 insertions(+), 41 deletions(-) diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 3b6806c73..49961536a 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -154,6 +154,59 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReaders( + std::vector>&& readers) { + auto raw_readers_guard = ScopeGuard([&]() -> void { + for (std::unique_ptr& reader : readers) { + if (reader != nullptr) { + reader->Close(); + } + } + }); + if (readers.empty()) { + return Status::Invalid("sorted readers must not be empty"); + } + for (const std::unique_ptr& reader : readers) { + if (reader == nullptr) { + return Status::Invalid("sorted readers must not contain null reader"); + } + } + + auto sort_merge_reader = std::make_unique( + std::move(readers), key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_); + auto create_consumer = [target_schema = write_schema_, pool = pool_]() + -> Result>> { + return KeyValueMetaProjectionConsumer::Create(target_schema, pool); + }; + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), 1, pool_); + raw_readers_guard.Release(); + ScopeGuard async_readers_guard([&]() -> void { async_key_value_producer_consumer->Close(); }); + std::unique_ptr>> rolling_writer; + PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); + ScopeGuard abort_writer_guard([&]() -> void { rolling_writer->Abort(); }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_key_value_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + } + PAIMON_RETURN_NOT_OK(rolling_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, + rolling_writer->GetResult()); + abort_writer_guard.Release(); + + for (const std::shared_ptr& flushed_file : flushed_files) { + new_files_.emplace_back(flushed_file); + PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); + } + metrics_->Merge(rolling_writer->GetMetrics()); + return Status::OK(); +} + Status MergeTreeWriter::Compact(bool full_compaction) { return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction); } @@ -256,49 +309,9 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, wait_for_latest_compaction = true; } auto cleanup_guard = ScopeGuard([&]() { write_buffer_->Clear(); }); - // 1. flush write buffer to get sorted readers PAIMON_ASSIGN_OR_RAISE(std::vector> readers, write_buffer_->CreateReaders()); - // 2. prepare loser tree sort merge reader - auto sort_merge_reader = std::make_unique( - std::move(readers), key_comparator_, user_defined_seq_comparator_, - merge_function_wrapper_); - // 3. project key value to arrow array - auto create_consumer = [target_schema = write_schema_, pool = pool_]() - -> Result>> { - return KeyValueMetaProjectionConsumer::Create(target_schema, pool); - }; - // consumer batch size is WriteBatchSize - auto async_key_value_producer_consumer = - std::make_unique>( - std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), - /*projection_thread_num=*/1, pool_); - std::unique_ptr>> - rolling_writer; - PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); - ScopeGuard write_guard([&]() -> void { - rolling_writer->Abort(); - async_key_value_producer_consumer->Close(); - }); - while (true) { - PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, - async_key_value_producer_consumer->NextBatch()); - if (key_value_batch.batch == nullptr) { - break; - } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); - } - PAIMON_RETURN_NOT_OK(rolling_writer->Close()); - PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, - rolling_writer->GetResult()); - async_key_value_producer_consumer->Close(); - write_guard.Release(); - - for (const auto& flushed_file : flushed_files) { - new_files_.emplace_back(flushed_file); - PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); - } - metrics_->Merge(rolling_writer->GetMetrics()); + PAIMON_RETURN_NOT_OK(WriteSortedReaders(std::move(readers))); } PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction)); PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction)); diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index febce2afb..542affd81 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -51,6 +51,7 @@ class IOManager; class FieldsComparator; class MemoryPool; class Metrics; +class KeyValueRecordReader; template class MergeFunctionWrapper; @@ -69,6 +70,8 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; + Status WriteSortedReaders(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 2155647a1..675ce3198 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -42,6 +43,7 @@ #include "paimon/core/io/compact_increment.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_increment.h" +#include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -52,6 +54,8 @@ #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/read_result_collector.h" @@ -64,6 +68,60 @@ class MergeFunctionWrapper; } // namespace paimon namespace paimon::test { +namespace { + +class TrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + TrackingKeyValueRecordReader(std::unique_ptr&& inner_reader, + bool* closed_flag) + : inner_reader_(std::move(inner_reader)), closed_flag_(closed_flag) {} + + Result> NextBatch() override { + return inner_reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return inner_reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_flag_ != nullptr) { + *closed_flag_ = true; + } + inner_reader_->Close(); + } + + private: + std::unique_ptr inner_reader_; + bool* closed_flag_; +}; + +class ErrorKeyValueRecordReader : public KeyValueRecordReader { + public: + ErrorKeyValueRecordReader(Status status, bool* closed_flag) + : status_(std::move(status)), closed_flag_(closed_flag) {} + + Result> NextBatch() override { + return status_; + } + + std::shared_ptr GetReaderMetrics() const override { + return nullptr; + } + + void Close() override { + if (closed_flag_ != nullptr) { + *closed_flag_ = true; + } + } + + private: + Status status_; + bool* closed_flag_; +}; + +} + class MergeTreeWriterTest : public ::testing::TestWithParam { public: class FakeCompactManager : public paimon::CompactManager { @@ -211,6 +269,21 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { writer_compact_manager, io_manager, /*enable_multi_thread_spill=*/false, pool_); } + std::unique_ptr CreateSingleReader( + const std::shared_ptr& array, int32_t batch_size = 16) const { + std::vector write_fields = {SpecialFields::SequenceNumber(), + SpecialFields::ValueKind()}; + write_fields.insert(write_fields.end(), value_fields_.begin(), value_fields_.end()); + std::shared_ptr write_schema = + DataField::ConvertDataFieldsToArrowSchema(write_fields); + std::shared_ptr key_schema = + arrow::schema(arrow::FieldVector({write_schema->field(2)})); + auto file_batch_reader = + std::make_unique(array, array->type(), batch_size); + return std::make_unique( + std::move(file_batch_reader), key_schema, value_schema_, 0, pool_); + } + private: std::shared_ptr pool_; std::shared_ptr file_system_; @@ -293,6 +366,29 @@ TEST_P(MergeTreeWriterTest, TestSimple) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [2, 0, "Alice", 10, 0, 13.1], + [0, 0, "Lucy", 20, 1, 14.1], + [1, 0, "Paul", 20, 1, null] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(-1, dir->Str(), sorted_reader_path_factory, 1, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); } TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { @@ -374,6 +470,146 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [16, 0, "Alice", 10, 0, 113.1], + [14, 0, "Lucy", 20, 1, 114.1], + [13, 0, "Paul", 20, 1, 15.1], + [15, 0, "Skye", 10, 0, 118.1] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(9, dir->Str(), sorted_reader_path_factory, 0, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + std::string uuid = path_factory->uuid_; + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(4, dir->Str(), path_factory, 7, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])") + .ValueOrDie()); + + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + std::string expected_data_file_name = "data-" + uuid + "-0.orc"; + std::string expected_data_file_path = dir->Str() + "/" + expected_data_file_name; + ASSERT_OK_AND_ASSIGN(FileStatus data_file_status, + options.GetFileSystem()->GetFileStatus(expected_data_file_path)); + + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])"}, + &expected_array) + .ok()); + CheckFileContent(expected_data_file_path, expected_array); + + ASSERT_TRUE(commit_increment.GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(expected_data_file_name, new_file->file_name); + ASSERT_EQ(data_file_status.GetLen(), new_file->file_size); + ASSERT_EQ(3, new_file->row_count); + ASSERT_EQ(7, new_file->min_sequence_number); + ASSERT_EQ(9, new_file->max_sequence_number); + ASSERT_EQ(7, new_file->schema_id); + ASSERT_EQ(1, new_file->delete_row_count); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderOwnership) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + + bool closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(closed); + ASSERT_OK(merge_writer->Close()); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + std::vector> empty_readers; + Status empty_status = merge_writer->WriteSortedReaders(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReaders(std::move(null_readers)); + ASSERT_TRUE(null_status.IsInvalid()); + + bool failing_reader_closed = false; + auto failing_reader = std::make_unique( + Status::IOError("sorted reader failure"), &failing_reader_closed); + std::vector> failing_readers; + failing_readers.push_back(std::move(failing_reader)); + Status failing_status = merge_writer->WriteSortedReaders(std::move(failing_readers)); + ASSERT_TRUE(failing_status.IsIOError()); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); } TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { From df322c1d38dc7e30bd2ca44c4d85297eddf96aac Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:22:00 +0800 Subject: [PATCH 14/24] feat(realtime): adapt prepared primary-key batches --- src/paimon/CMakeLists.txt | 1 + .../merged_key_value_record_reader_test.cpp | 398 ++++++++++++ .../core/io/prepared_key_value_reader.cpp | 565 ++++++++++++++++++ .../core/io/prepared_key_value_reader.h | 41 ++ src/paimon/core/realtime/realtime_fields.h | 37 ++ .../core/schema/schema_validation_test.cpp | 7 + 6 files changed, 1049 insertions(+) create mode 100644 src/paimon/core/io/prepared_key_value_reader.cpp create mode 100644 src/paimon/core/io/prepared_key_value_reader.h create mode 100644 src/paimon/core/realtime/realtime_fields.h diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index b2e1e6617..fc1fb00dd 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -282,6 +282,7 @@ set(PAIMON_CORE_SRCS core/io/vector_file_batch_reader.cpp core/io/file_index_evaluator.cpp core/io/key_value_data_file_record_reader.cpp + core/io/prepared_key_value_reader.cpp core/io/key_value_data_file_writer_factory.cpp core/io/key_value_data_file_writer.cpp core/io/key_value_in_memory_record_reader.cpp diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 1b6b71c69..39714fa29 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -18,8 +18,13 @@ #include "paimon/core/io/merged_key_value_record_reader.h" +#include +#include #include +#include +#include #include +#include #include "arrow/api.h" #include "arrow/array/array_nested.h" @@ -27,10 +32,14 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/io/prepared_key_value_reader.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/realtime/realtime_fields.h" #include "paimon/memory/memory_pool.h" +#include "paimon/realtime/offset_range.h" #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/key_value_checker.h" @@ -38,6 +47,56 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { + +namespace { + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))); +} + +std::shared_ptr MakePreparedSchema(const arrow::FieldVector& value_fields) { + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), value_fields.begin(), value_fields.end()); + return arrow::schema(prepared_fields); +} + +class TrackingBatchReader : public BatchReader { + public: + TrackingBatchReader(std::unique_ptr&& delegate, int32_t* close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + if (closed_) { + return; + } + closed_ = true; + ++(*close_count_); + delegate_->Close(); + } + + private: + bool closed_ = false; + std::unique_ptr delegate_; + int32_t* close_count_; +}; + +} + class MergedKeyValueRecordReaderTest : public testing::Test { public: void SetUp() override { @@ -51,6 +110,14 @@ class MergedKeyValueRecordReaderTest : public testing::Test { std::shared_ptr merge_function_wrapper_; }; +TEST_F(MergedKeyValueRecordReaderTest, TestRealtimeOffsetField) { + const DataField& field = RealtimeOffsetField(); + ASSERT_EQ(std::numeric_limits::max() - 10002, field.Id()); + ASSERT_EQ("_REALTIME_OFFSET", field.Name()); + ASSERT_EQ(arrow::Type::INT64, field.Type()->id()); + ASSERT_FALSE(field.Nullable()); +} + TEST_F(MergedKeyValueRecordReaderTest, TestMergeAcrossUnderlyingBatches) { std::vector fields = {DataField(0, arrow::field("k0", arrow::int32())), DataField(1, arrow::field("k1", arrow::int32())), @@ -143,4 +210,335 @@ TEST_F(MergedKeyValueRecordReaderTest, TestSkipMergedNulloptResultInHasNext) { } } +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderOffsetFilter) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 100, 0, 1, 10], + [0, 101, 1, 2, 20], + [0, 102, 4, 3, 30], + [0, 103, 2, 4, 40], + [0, 104, 5, 5, 50], + [0, 105, 3, 6, 60] + ])") + .ValueOrDie()); + + auto batch_reader = + std::make_unique(prepared_array, prepared_type, 2); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), + key_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + std::vector row_kinds = {const_cast(RowKind::Insert()), + const_cast(RowKind::Insert())}; + std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; + std::vector expected = KeyValueChecker::GenerateKeyValues( + row_kinds, {103, 105}, levels, {{4}, {6}}, {{4, 40}, {6, 60}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderRawCountBeforeDedup) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1, 100], + [2, 11, 1, 1, 101], + [0, 12, 2, 2, 200] + ])") + .ValueOrDie()); + + int64_t raw_row_count = 0; + auto batch_reader = + std::make_unique(prepared_array, prepared_type, 2); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr raw_reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, std::nullopt, key_schema, + value_schema, pool_, &raw_row_count)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({value_fields[0]}, true)); + auto merged_reader = std::make_unique( + std::move(raw_reader), key_comparator, merge_function_wrapper_); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult< + MergedKeyValueRecordReader, KeyValueRecordReader::Iterator>(merged_reader.get()))); + + ASSERT_EQ(raw_row_count, 3); + std::vector row_kinds = {const_cast(RowKind::UpdateAfter()), + const_cast(RowKind::Insert())}; + std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; + std::vector expected = KeyValueChecker::GenerateKeyValues( + row_kinds, {11, 12}, levels, {{1}, {2}}, {{1, 101}, {2, 200}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderRawCountBeforeFilter) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1], + [0, 11, 1, 2], + [0, 12, 2, 3], + [0, 13, 3, 4] + ])") + .ValueOrDie()); + + int64_t raw_row_count = 0; + auto batch_reader = + std::make_unique(prepared_array, prepared_type, 2); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), + value_schema, value_schema, pool_, &raw_row_count)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + ASSERT_EQ(results.size(), 2); + ASSERT_EQ(raw_row_count, 4); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr extra = MakeField("extra", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, extra}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1, 2]])") + .ValueOrDie()); + + auto query_batch_reader = + std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr query_reader, + AdaptPreparedBatchReader(std::move(query_batch_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector query_results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(query_reader.get()))); + ASSERT_EQ(query_results.size(), 1); + ASSERT_EQ(query_results[0].value->GetFieldCount(), 1); + ASSERT_EQ(query_results[0].value->GetInt(0), 1); + + auto batch_reader = + std::make_unique(prepared_array, prepared_type, 1); + ASSERT_NOK_WITH_MSG(AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, + std::nullopt, value_schema, value_schema, pool_), + "exact"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderSafeDecode) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + + arrow::FieldVector invalid_fields = prepared_schema->fields(); + invalid_fields[0] = invalid_fields[0]->WithName("wrong_value_kind"); + invalid_fields[3] = MakeField("wrong_key", arrow::int32(), 99); + std::shared_ptr invalid_type = arrow::struct_(invalid_fields); + auto invalid_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[0, 10, 0, 1]])").ValueOrDie()); + + auto batch_reader = + std::make_unique(invalid_array, invalid_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "prepared batch field"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { + std::shared_ptr id = MakeField("id", arrow::int32(), 0); + std::shared_ptr item_a = MakeField("a", arrow::int32(), 10); + std::shared_ptr item_b = MakeField("b", arrow::int32(), 11); + std::shared_ptr items = + MakeField("items", arrow::list(arrow::field("item", arrow::struct_({item_a, item_b}))), 2); + std::shared_ptr attr_x = MakeField("x", arrow::int32(), 20); + std::shared_ptr attr_y = MakeField("y", arrow::int32(), 21); + std::shared_ptr attrs = + MakeField("attrs", arrow::map(arrow::utf8(), arrow::struct_({attr_x, attr_y})), 3); + std::shared_ptr key_left = MakeField("left", arrow::int32(), 30); + std::shared_ptr key_right = MakeField("right", arrow::int32(), 31); + std::shared_ptr keyed_values = MakeField( + "keyed_values", arrow::map(arrow::struct_({key_left, key_right}), arrow::int32()), 4); + std::shared_ptr full_value_schema = + arrow::schema({id, items, attrs, keyed_values}); + std::shared_ptr key_schema = arrow::schema({id}); + std::shared_ptr prepared_schema = + MakePreparedSchema(full_value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1, [[100, 200], [300, 400]], [["k1", [7, 8]], ["k2", [9, 10]]], [[[11, 12], 13], [[21, 22], 23]]] + ])") + .ValueOrDie()); + + std::shared_ptr query_item_b = MakeField("renamed_b", arrow::int32(), 11); + std::shared_ptr query_item_a = MakeField("renamed_a", arrow::int32(), 10); + std::shared_ptr query_items = MakeField( + "items_renamed", + arrow::list(arrow::field("element", arrow::struct_({query_item_b, query_item_a}))), 2); + std::shared_ptr query_attr_y = MakeField("renamed_y", arrow::int32(), 21); + std::shared_ptr query_attr_x = MakeField("renamed_x", arrow::int32(), 20); + std::shared_ptr query_attrs = + MakeField("attrs_renamed", + arrow::map(arrow::utf8(), arrow::struct_({query_attr_y, query_attr_x})), 3); + std::shared_ptr query_key_right = MakeField("renamed_right", arrow::int32(), 31); + std::shared_ptr query_key_left = MakeField("renamed_left", arrow::int32(), 30); + std::shared_ptr query_keyed_values = + MakeField("keyed_values_renamed", + arrow::map(arrow::struct_({query_key_right, query_key_left}), arrow::int32()), 4); + std::shared_ptr query_value_schema = + arrow::schema({id, query_items, query_attrs, query_keyed_values}); + + auto batch_reader = + std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + key_schema, query_value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0].key->GetInt(0), 1); + ASSERT_EQ(results[0].value->GetFieldCount(), 4); + ASSERT_EQ(results[0].value->GetInt(0), 1); + + std::shared_ptr item_array = results[0].value->GetArray(1); + ASSERT_EQ(item_array->Size(), 2); + std::shared_ptr first_item = item_array->GetRow(0, 2); + ASSERT_EQ(first_item->GetInt(0), 200); + ASSERT_EQ(first_item->GetInt(1), 100); + std::shared_ptr second_item = item_array->GetRow(1, 2); + ASSERT_EQ(second_item->GetInt(0), 400); + ASSERT_EQ(second_item->GetInt(1), 300); + + std::shared_ptr attr_map = results[0].value->GetMap(2); + ASSERT_EQ(attr_map->Size(), 2); + std::shared_ptr key_array = attr_map->KeyArray(); + ASSERT_EQ(std::string(key_array->GetStringView(0)), "k1"); + ASSERT_EQ(std::string(key_array->GetStringView(1)), "k2"); + std::shared_ptr value_array = attr_map->ValueArray(); + std::shared_ptr first_attr = value_array->GetRow(0, 2); + ASSERT_EQ(first_attr->GetInt(0), 8); + ASSERT_EQ(first_attr->GetInt(1), 7); + std::shared_ptr second_attr = value_array->GetRow(1, 2); + ASSERT_EQ(second_attr->GetInt(0), 10); + ASSERT_EQ(second_attr->GetInt(1), 9); + + std::shared_ptr keyed_value_map = results[0].value->GetMap(3); + ASSERT_EQ(keyed_value_map->Size(), 2); + std::shared_ptr struct_keys = keyed_value_map->KeyArray(); + std::shared_ptr first_key = struct_keys->GetRow(0, 2); + ASSERT_EQ(first_key->GetInt(0), 12); + ASSERT_EQ(first_key->GetInt(1), 11); + std::shared_ptr second_key = struct_keys->GetRow(1, 2); + ASSERT_EQ(second_key->GetInt(0), 22); + ASSERT_EQ(second_key->GetInt(1), 21); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(0), 13); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(1), 23); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderClose) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1, 100] + ])") + .ValueOrDie()); + + int32_t explicit_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &explicit_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + reader->Close(); + reader->Close(); + } + ASSERT_EQ(explicit_close_count, 1); + + int32_t destructor_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &destructor_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + } + ASSERT_EQ(destructor_close_count, 1); + + int32_t factory_failure_close_count = 0; + { + std::unique_ptr tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &factory_failure_close_count); + std::shared_ptr invalid_schema = arrow::schema(value_schema->fields()); + ASSERT_NOK(AdaptPreparedBatchReader(std::move(tracking_reader), invalid_schema, + OffsetRange(0, 1), key_schema, value_schema, pool_)); + ASSERT_EQ(nullptr, tracking_reader); + } + ASSERT_EQ(factory_failure_close_count, 1); + + int32_t read_failure_close_count = 0; + { + auto failing_reader = + std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("prepared reader failure")); + auto tracking_reader = std::make_unique(std::move(failing_reader), + &read_failure_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "prepared reader failure"); + ASSERT_EQ(read_failure_close_count, 1); + reader->Close(); + } + ASSERT_EQ(read_failure_close_count, 1); +} + } // namespace paimon::test diff --git a/src/paimon/core/io/prepared_key_value_reader.cpp b/src/paimon/core/io/prepared_key_value_reader.cpp new file mode 100644 index 000000000..0f4f22097 --- /dev/null +++ b/src/paimon/core/io/prepared_key_value_reader.cpp @@ -0,0 +1,565 @@ +/* + * 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/io/prepared_key_value_reader.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_nested.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.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/scope_guard.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/status.h" + +namespace paimon { + +namespace { + +constexpr int32_t kValueKindIndex = 0; +constexpr int32_t kSequenceNumberIndex = 1; +constexpr int32_t kRealtimeOffsetIndex = 2; +constexpr int32_t kPreparedValueStartIndex = 3; + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type); + +Status CheckPreparedField(const std::shared_ptr& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid(fmt::format("prepared schema missing transport field {} at index {}", + expected_field.Name(), field_idx)); + } + const std::shared_ptr& field = schema->field(field_idx); + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field->name() != expected_field.Name() || !field->type()->Equals(*expected_field.Type()) || + field->nullable() || field_id != expected_field.Id()) { + return Status::Invalid(fmt::format( + "prepared schema field {} must be non-null {}:{} with field id {}, got {}:{} " + "nullable={} field id {}", + field_idx, expected_field.Name(), expected_field.Type()->ToString(), + expected_field.Id(), field->name(), field->type()->ToString(), field->nullable(), + field_id)); + } + return Status::OK(); +} + +Result FindFieldIndexByPaimonId(const arrow::FieldVector& fields, int32_t field_id) { + std::optional matching_index; + for (int32_t i = 0; i < static_cast(fields.size()); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t candidate_id, + NestedProjectionUtils::GetPaimonFieldId(fields[i])); + if (candidate_id == field_id) { + if (matching_index.has_value()) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + matching_index = i; + } + } + if (matching_index.has_value()) { + return matching_index.value(); + } + return Status::Invalid(fmt::format("cannot find field id {} in prepared schema", field_id)); +} + +Status ValidateProjectionType(const std::shared_ptr& prepared_type, + const std::shared_ptr& query_type) { + if (prepared_type->id() != query_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + switch (query_type->id()) { + case arrow::Type::STRUCT: { + const arrow::FieldVector& prepared_fields = prepared_type->fields(); + for (const std::shared_ptr& query_field : query_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); + } + case arrow::Type::LIST: + return ValidateProjectionType(prepared_type->field(0)->type(), + query_type->field(0)->type()); + case arrow::Type::MAP: { + const std::shared_ptr prepared_map = + checked_pointer_cast(prepared_type); + const std::shared_ptr query_map = + checked_pointer_cast(query_type); + PAIMON_RETURN_NOT_OK( + ValidateProjectionType(prepared_map->key_type(), query_map->key_type())); + return ValidateProjectionType(prepared_map->item_type(), query_map->item_type()); + } + default: + if (!prepared_type->Equals(*query_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + return Status::OK(); + } +} + +Status ValidateProjectionSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema) { + arrow::FieldVector prepared_value_fields( + prepared_schema->fields().begin() + kPreparedValueStartIndex, + prepared_schema->fields().end()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_value_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_value_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); +} + +Status ValidateExactCommitSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& value_schema) { + if (prepared_schema->num_fields() != value_schema->num_fields() + kPreparedValueStartIndex) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + for (int32_t i = 0; i < value_schema->num_fields(); ++i) { + if (!prepared_schema->field(i + kPreparedValueStartIndex) + ->Equals(value_schema->field(i), true)) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + } + return Status::OK(); +} + +Status ValidatePreparedSchema(const std::shared_ptr& prepared_schema) { + if (!prepared_schema || prepared_schema->num_fields() < kPreparedValueStartIndex) { + return Status::Invalid("prepared schema must contain realtime transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kSequenceNumberIndex, SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kRealtimeOffsetIndex, RealtimeOffsetField())); + return Status::OK(); +} + +Result> AlignStructArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type) { + const std::shared_ptr data_type = + checked_pointer_cast(array->type()); + std::unordered_map data_field_id_to_idx; + data_field_id_to_idx.reserve(data_type->num_fields()); + for (int32_t i = 0; i < data_type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(data_type->field(i))); + if (!data_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared value struct", field_id)); + } + } + + arrow::ArrayVector aligned_arrays; + aligned_arrays.reserve(read_type->num_fields()); + for (const std::shared_ptr& read_field : read_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t read_field_id, + NestedProjectionUtils::GetPaimonFieldId(read_field)); + auto data_iter = data_field_id_to_idx.find(read_field_id); + if (data_iter == data_field_id_to_idx.end()) { + return Status::Invalid( + fmt::format("cannot find field id {} in prepared value struct", read_field_id)); + } + std::shared_ptr child = array->field(data_iter->second); + PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type())); + aligned_arrays.push_back(std::move(child)); + } + + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr aligned, + arrow::StructArray::Make(aligned_arrays, read_type->fields(), array->null_bitmap(), + array->null_count(), array->offset())); + return aligned; +} + +Result> AlignListArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type) { + std::shared_ptr values = array->values(); + PAIMON_ASSIGN_OR_RAISE(values, AlignArrayByPaimonIds(values, read_type->value_type())); + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {values->data()}; + return arrow::MakeArray(new_data); +} + +Result> AlignMapArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type) { + std::shared_ptr keys = array->keys(); + PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type())); + std::shared_ptr items = array->items(); + PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type())); + + const std::shared_ptr& entries_data = array->data()->child_data[0]; + std::shared_ptr new_entries = entries_data->Copy(); + new_entries->type = arrow::struct_({read_type->key_field(), read_type->item_field()}); + new_entries->child_data = {keys->data(), items->data()}; + + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {std::move(new_entries)}; + return arrow::MakeArray(new_data); +} + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type) { + if (array->type()->id() != read_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + switch (read_type->id()) { + case arrow::Type::STRUCT: + return AlignStructArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type)); + case arrow::Type::LIST: + return AlignListArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type)); + case arrow::Type::MAP: + return AlignMapArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type)); + default: + if (!array->type()->Equals(*read_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + return array; + } +} + +Result ProjectFieldsByPaimonIds( + const std::shared_ptr& data_batch, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema) { + std::unordered_map prepared_field_id_to_idx; + prepared_field_id_to_idx.reserve(prepared_schema->num_fields()); + for (int32_t i = kPreparedValueStartIndex; i < prepared_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(prepared_schema->field(i))); + if (!prepared_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + } + + arrow::ArrayVector result; + result.reserve(query_schema->num_fields()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_field_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + auto prepared_iter = prepared_field_id_to_idx.find(query_field_id); + if (prepared_iter == prepared_field_id_to_idx.end()) { + return Status::Invalid( + fmt::format("cannot find field id {} in prepared schema", query_field_id)); + } + std::shared_ptr field_array = data_batch->field(prepared_iter->second); + PAIMON_ASSIGN_OR_RAISE(field_array, + AlignArrayByPaimonIds(field_array, query_field->type())); + result.push_back(std::move(field_array)); + } + return result; +} + +Result> ApplyOffsetFilter( + const std::shared_ptr& data_batch, + const std::shared_ptr>& offset_array, + const std::optional& visible_offsets, arrow::MemoryPool* arrow_pool) { + if (!visible_offsets.has_value()) { + return data_batch; + } + + arrow::BooleanBuilder filter_builder(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(filter_builder.Reserve(offset_array->length())); + int64_t visible_row_count = 0; + for (int64_t i = 0; i < offset_array->length(); ++i) { + int64_t offset = offset_array->Value(i); + bool visible = offset >= visible_offsets->begin && offset < visible_offsets->end; + filter_builder.UnsafeAppend(visible); + visible_row_count += visible; + } + if (visible_row_count == 0) { + return std::shared_ptr(); + } + if (visible_row_count == data_batch->length()) { + return data_batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr filter, + filter_builder.Finish()); + arrow::compute::ExecContext exec_context(arrow_pool); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum filtered, + arrow::compute::Filter(data_batch, filter, arrow::compute::FilterOptions::Defaults(), + &exec_context)); + return checked_pointer_cast(filtered.make_array()); +} + +class PreparedKeyValueReader final : public KeyValueRecordReader { + public: + PreparedKeyValueReader(std::unique_ptr&& reader, + const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& pool, int64_t* raw_row_count) + : reader_(std::move(reader)), + prepared_schema_(prepared_schema), + visible_offsets_(visible_offsets), + key_schema_(key_schema), + value_schema_(value_schema), + pool_(pool), + arrow_pool_(GetArrowPool(pool)), + raw_row_count_(raw_row_count) {} + + ~PreparedKeyValueReader() override { + Close(); + } + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(PreparedKeyValueReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->row_kind_array_->length(); + } + + Result Next() override { + if (cursor_ >= reader_->row_kind_array_->length()) { + return Status::Invalid("No more prepared key values in current iterator"); + } + std::shared_ptr key = + std::make_shared(reader_->key_ctx_, cursor_); + auto value = std::make_unique(reader_->value_ctx_, cursor_); + PAIMON_ASSIGN_OR_RAISE( + const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kind_array_->Value(cursor_))); + int64_t sequence_number = reader_->sequence_number_array_->Value(cursor_); + ++cursor_; + return KeyValue(row_kind, sequence_number, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + PreparedKeyValueReader* reader_; + int64_t cursor_ = 0; + }; + + Result> NextBatch() override { + Result> result = NextBatchImpl(); + if (!result.ok()) { + Close(); + } + return result; + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_) { + return; + } + closed_ = true; + ResetBatchState(); + reader_->Close(); + } + + private: + Result> NextBatchImpl() { + if (closed_) { + return std::unique_ptr(); + } + + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return std::unique_ptr(); + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("cannot cast prepared batch to StructArray"); + } + std::shared_ptr data_batch = + checked_pointer_cast(arrow_array); + PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch)); + if (raw_row_count_ != nullptr) { + int64_t updated_count = 0; + if (__builtin_add_overflow(*raw_row_count_, data_batch->length(), &updated_count)) { + return Status::Invalid("prepared raw row count overflow"); + } + *raw_row_count_ = updated_count; + } + + std::shared_ptr> offset_array = + checked_pointer_cast>( + data_batch->field(kRealtimeOffsetIndex)); + PAIMON_ASSIGN_OR_RAISE( + data_batch, + ApplyOffsetFilter(data_batch, offset_array, visible_offsets_, arrow_pool_.get())); + if (!data_batch) { + continue; + } + + row_kind_array_ = checked_pointer_cast>( + data_batch->field(kValueKindIndex)); + sequence_number_array_ = checked_pointer_cast>( + data_batch->field(kSequenceNumberIndex)); + PAIMON_ASSIGN_OR_RAISE( + arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_)); + PAIMON_ASSIGN_OR_RAISE( + arrow::ArrayVector value_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, value_schema_)); + key_ctx_ = std::make_shared(key_fields, pool_); + value_ctx_ = std::make_shared(value_fields, pool_); + ArrowUtils::TraverseArray(data_batch); + return std::make_unique(this); + } + } + + Status ValidatePreparedBatch(const std::shared_ptr& data_batch) const { + if (data_batch->num_fields() != prepared_schema_->num_fields()) { + return Status::Invalid(fmt::format( + "prepared batch field count {} does not match prepared schema field count {}", + data_batch->num_fields(), prepared_schema_->num_fields())); + } + const arrow::FieldVector& batch_fields = data_batch->type()->fields(); + for (int32_t i = 0; i < data_batch->num_fields(); ++i) { + if (!batch_fields[i]->Equals(prepared_schema_->field(i), true)) { + return Status::Invalid(fmt::format( + "prepared batch field {} does not match declared prepared schema", i)); + } + } + if (!data_batch->field(kValueKindIndex) || + data_batch->field(kValueKindIndex)->type_id() != arrow::Type::INT8) { + return Status::Invalid("cannot cast VALUE_KIND column to int8 arrow array"); + } + if (!data_batch->field(kSequenceNumberIndex) || + data_batch->field(kSequenceNumberIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast SEQUENCE_NUMBER column to int64 arrow array"); + } + if (!data_batch->field(kRealtimeOffsetIndex) || + data_batch->field(kRealtimeOffsetIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast REALTIME_OFFSET column to int64 arrow array"); + } + if (data_batch->field(kValueKindIndex)->null_count() != 0 || + data_batch->field(kSequenceNumberIndex)->null_count() != 0 || + data_batch->field(kRealtimeOffsetIndex)->null_count() != 0) { + return Status::Invalid("prepared transport columns must not contain nulls"); + } + return Status::OK(); + } + + void ResetBatchState() { + key_ctx_.reset(); + value_ctx_.reset(); + row_kind_array_.reset(); + sequence_number_array_.reset(); + } + + private: + bool closed_ = false; + std::unique_ptr reader_; + std::shared_ptr prepared_schema_; + std::optional visible_offsets_; + std::shared_ptr key_schema_; + std::shared_ptr value_schema_; + std::shared_ptr pool_; + std::shared_ptr arrow_pool_; + int64_t* raw_row_count_; + std::shared_ptr key_ctx_; + std::shared_ptr value_ctx_; + std::shared_ptr> row_kind_array_; + std::shared_ptr> sequence_number_array_; +}; + +} + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool, int64_t* raw_row_count) { + std::unique_ptr owned_reader = std::move(reader); + if (!owned_reader) { + return Status::Invalid("prepared batch reader cannot be null"); + } + ScopeGuard close_guard([&owned_reader]() -> void { owned_reader->Close(); }); + PAIMON_RETURN_NOT_OK(ValidatePreparedSchema(prepared_schema)); + if (!key_schema) { + return Status::Invalid("prepared key schema cannot be null"); + } + if (!value_schema) { + return Status::Invalid("prepared value schema cannot be null"); + } + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, key_schema)); + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, value_schema)); + if (!visible_offsets.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateExactCommitSchema(prepared_schema, value_schema)); + } + std::unique_ptr result( + new PreparedKeyValueReader(std::move(owned_reader), prepared_schema, visible_offsets, + key_schema, value_schema, memory_pool, raw_row_count)); + close_guard.Release(); + return result; +} + +} diff --git a/src/paimon/core/io/prepared_key_value_reader.h b/src/paimon/core/io/prepared_key_value_reader.h new file mode 100644 index 000000000..e7a6f9651 --- /dev/null +++ b/src/paimon/core/io/prepared_key_value_reader.h @@ -0,0 +1,41 @@ +/* + * 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 "arrow/type_fwd.h" +#include "paimon/core/io/key_value_record_reader.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/result.h" + +namespace paimon { +class BatchReader; +class MemoryPool; + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool, int64_t* raw_row_count = nullptr); + +} diff --git a/src/paimon/core/realtime/realtime_fields.h b/src/paimon/core/realtime/realtime_fields.h new file mode 100644 index 000000000..6ed04b38a --- /dev/null +++ b/src/paimon/core/realtime/realtime_fields.h @@ -0,0 +1,37 @@ +/* + * 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 "arrow/type.h" +#include "paimon/common/types/data_field.h" + +namespace paimon { + +inline const DataField& RealtimeOffsetField() { + static const DataField data_field = + DataField(std::numeric_limits::max() - 10002, + arrow::field("_REALTIME_OFFSET", arrow::int64(), /*nullable=*/false)); + return data_field; +} + +} // namespace paimon diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 47603497b..050f09701 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -46,6 +46,13 @@ TEST(SchemaValidationTest, TestSimple) { ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestRealtimeOffsetIsNotGloballyReserved) { + auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(0, schema, {}, {}, {})); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); +} + TEST(SchemaValidationTest, TestVectorType) { auto vector_field = arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)); auto schema = arrow::schema({arrow::field("id", arrow::int64()), vector_field}); From 87e454803670fa909fb5a711bca657d2375918fa Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:26:17 +0800 Subject: [PATCH 15/24] refactor(realtime): prepare primary-key batches in framework --- include/paimon/realtime/realtime_context.h | 4 + include/paimon/realtime/realtime_store.h | 52 +- .../core/mergetree/merge_tree_writer_test.cpp | 15 + src/paimon/core/mergetree/write_buffer.cpp | 4 + .../operation/key_value_file_store_write.cpp | 40 +- .../key_value_file_store_write_test.cpp | 374 +++++++++++- .../core/operation/merge_file_split_read.cpp | 10 +- .../core/operation/merge_file_split_read.h | 1 - .../realtime/arrow_realtime_store_factory.cpp | 32 +- .../realtime/primary_key_realtime_store.cpp | 548 ++++-------------- .../realtime/primary_key_realtime_store.h | 27 +- .../primary_key_realtime_store_test.cpp | 502 +++------------- .../core/realtime/realtime_context_impl.cpp | 41 +- .../core/realtime/realtime_context_impl.h | 5 - .../core/realtime/realtime_context_test.cpp | 133 ++--- .../realtime/realtime_primary_key_writer.cpp | 332 +++++++---- .../realtime/realtime_primary_key_writer.h | 35 +- .../table/source/key_value_table_read.cpp | 193 ++---- test/inte/realtime_write_inte_test.cpp | 239 ++++---- 19 files changed, 1106 insertions(+), 1481 deletions(-) diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index 200e4ba4c..8f2967b32 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -78,6 +78,10 @@ using RealtimeOffsetMap = std::map; /// partition drop, and rollback operations do not automatically clear process-local real-time /// state. Applications must coordinate these operations with active real-time writers and recreate /// the `RealtimeContext` and writers before continuing. +/// +/// A primary-key writer and its context form one lifecycle. After a primary-key write or prepare +/// returns an error, discard both, create fresh instances from the latest committed snapshot, and +/// replay the caller-owned WAL. Reusing the failed context or retrying prepare is unsupported. class PAIMON_EXPORT RealtimeContext { public: /// Creates a context backed by Paimon's default in-memory Arrow `RealtimeStore`. diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index 1e53c173e..dc5d543ac 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -47,19 +47,15 @@ struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig { StatisticsMode statistics_mode; }; -struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig { - std::vector primary_keys; - /// Largest sequence restored from the committed snapshot. A PK store assigns one contiguous - /// sequence to every mutation in `Write` order, starting at the next value, and rejects - /// `Write` before the assigned sequence would exceed `INT64_MAX - 1`. - int64_t restore_max_sequence_number; -}; +struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig {}; using RealtimeStoreCreateConfig = std::variant; struct PAIMON_EXPORT RealtimeStoreCreateRequest { - /// Complete table write schema whose ownership is transferred to the factory. + /// Schema whose ownership is transferred to the factory. Append mode receives the complete + /// table write schema. Primary-key mode receives the prepared transport schema: + /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields]. std::unique_ptr<::ArrowSchema> write_schema; std::map options; std::shared_ptr memory_pool; @@ -68,10 +64,12 @@ struct PAIMON_EXPORT RealtimeStoreCreateRequest { RealtimeStoreCreateConfig mode_config; }; -/// A table record batch and its framework-assigned contiguous offset range. +/// A record batch and its framework-assigned contiguous offset range. /// -/// The batch contains only table write fields. Row `i` is associated with -/// `offset_range.begin + i`; the offset is progress metadata and is not a table field. +/// Append-mode batches contain table write fields, and row `i` is associated with +/// `offset_range.begin + i`. Primary-key batches contain the prepared transport schema supplied +/// to the factory and are physically sorted by full primary key then sequence number; their +/// per-row `_REALTIME_OFFSET` field preserves the original write-order offset after sorting. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -106,7 +104,8 @@ class PAIMON_EXPORT RealtimeReadView { /// Parameters used by a `RealtimeStore` to create readers for a query. struct PAIMON_EXPORT RealtimeQueryContext { - /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added. + /// Append mode receives the requested output fields before the mandatory leading + /// `_VALUE_KIND` field is added. Primary-key mode receives the complete prepared schema. ::ArrowSchema* read_schema; /// Predicate using field indexes from `read_schema`. std::shared_ptr predicate; @@ -143,9 +142,13 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// - /// Concatenating the returned readers must produce every sealed row exactly once and in write - /// order. Each output batch contains `_VALUE_KIND` followed by all fields from the factory's - /// `write_schema`. + /// Concatenating the returned readers must produce every sealed row exactly once. Append-mode + /// readers preserve write order and contain `_VALUE_KIND` followed by table write fields. + /// Primary-key readers expose raw prepared rows. Each returned reader's complete stream, + /// including across `NextBatch` boundaries, is sorted by full primary key then sequence + /// number; all readers collectively cover sealed mutations exactly once. Reader cardinality is + /// independent of the number of writes. Paimon adapts and merges those rows before writing + /// files. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -155,16 +158,17 @@ class PAIMON_EXPORT RealtimeStore { /// also provide a consistent snapshot when a write or seal is in progress. virtual Result> AcquireReadView() = 0; - /// Creates readers over rows in `view` whose offsets are greater than or equal to - /// `offset_begin`. + /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater + /// than or equal to `offset_begin`. Primary-key mode ignores `offset_begin` and returns raw + /// prepared rows; Paimon applies offset filtering, projection, and merge-on-read adaptation. /// - /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by - /// `context.read_schema` except a duplicate `_VALUE_KIND`. Concatenating all returned readers - /// must produce every matching row once. Primary-key readers additionally provide a non-null - /// `_SEQUENCE_NUMBER` when requested, are individually sorted by primary key, and contain at - /// most one mutation per key. Assigned sequences remain stable across views and queries; - /// readers need not be globally sorted with one another. Paimon retains `view` for the lifetime - /// of the resulting framework reader. + /// Append-mode output batches contain `_VALUE_KIND` first, followed by requested fields except + /// a duplicate `_VALUE_KIND`; all returned append readers collectively cover every matching + /// row exactly once. Primary-key output batches use the prepared transport schema and may + /// contain multiple mutations per key. Each returned primary-key reader's complete stream is + /// sorted by full primary key then sequence number, and all readers collectively cover raw + /// mutations exactly once. Reader cardinality is independent of the number of writes. Paimon + /// retains `view` for the lifetime of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 675ce3198..aa2d0c959 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -612,6 +613,20 @@ TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { ASSERT_OK(merge_writer->Close()); } +TEST_P(MergeTreeWriterTest, TestRejectsExhaustedSequence) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_NOK_WITH_MSG(CreateMergeWriter(std::numeric_limits::max(), dir->Str(), + path_factory, 0, options), + "sequence number has reached INT64_MAX"); +} + TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({ diff --git a/src/paimon/core/mergetree/write_buffer.cpp b/src/paimon/core/mergetree/write_buffer.cpp index 549975a33..3d3fdc196 100644 --- a/src/paimon/core/mergetree/write_buffer.cpp +++ b/src/paimon/core/mergetree/write_buffer.cpp @@ -18,6 +18,7 @@ #include "paimon/core/mergetree/write_buffer.h" +#include #include #include @@ -39,6 +40,9 @@ Result> WriteBuffer::Create( const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, const std::shared_ptr& io_manager, bool enable_multi_thread_spill, const std::shared_ptr& pool) { + if (last_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("sequence number has reached INT64_MAX"); + } auto value_type = arrow::struct_(value_schema->fields()); auto in_memory_buffer = std::make_unique( last_sequence_number, value_type, trimmed_primary_keys, user_defined_sequence_fields, 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 492161cf8..d2c97abcf 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -18,12 +18,13 @@ #include "paimon/core/operation/key_value_file_store_write.h" -#include #include #include #include "arrow/c/bridge.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" @@ -35,6 +36,7 @@ #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" #include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_fields.h" #include "paimon/core/realtime/realtime_primary_key_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/utils/file_store_path_factory.h" @@ -124,7 +126,6 @@ Result> KeyValueFileStoreWrite::CreateWriter( std::shared_ptr levels, Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); std::map partition_map; - int64_t initial_max_sequence_number = restore_max_seq_number; std::shared_ptr compact_manager; std::shared_ptr realtime_context_impl; std::optional realtime_store_state; @@ -135,19 +136,27 @@ Result> KeyValueFileStoreWrite::CreateWriter( partition_map = std::map(partition_values.begin(), partition_values.end()); PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); + if (schema_->GetFieldByName(RealtimeOffsetField().Name())) { + return Status::Invalid("PK real-time write schema contains reserved transport field " + + RealtimeOffsetField().Name()); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), schema_->fields().begin(), + schema_->fields().end()); auto c_write_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, c_write_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(std::move(prepared_fields)), c_write_schema.get())); PAIMON_ASSIGN_OR_RAISE( RealtimeStoreState store_state, realtime_context_impl->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ std::move(c_write_schema), options_.ToMap(), pool_, partition_map, bucket, - PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys, - restore_max_seq_number}})); + PrimaryKeyRealtimeStoreCreateConfig{}})); realtime_store_state = std::move(store_state); - initial_max_sequence_number = realtime_store_state->initial_max_sequence_number.value(); - if (initial_max_sequence_number == std::numeric_limits::max()) { - return Status::Invalid("PK sequence number has reached INT64_MAX"); - } compact_manager = std::make_shared(); } else { auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); @@ -159,15 +168,16 @@ Result> KeyValueFileStoreWrite::CreateWriter( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( - initial_max_sequence_number, trimmed_primary_keys, data_file_path_factory, - key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, - table_schema_->Id(), schema_, options_, compact_manager, - realtime_context_ ? nullptr : io_manager_, enable_multi_thread_spill_, pool_)); + restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, + user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, + options_, compact_manager, realtime_context_ ? nullptr : io_manager_, + enable_multi_thread_spill_, pool_)); if (!realtime_context_) { return std::shared_ptr(std::move(writer)); } - return RealtimePrimaryKeyWriter::Create(partition_map, bucket, schema_, realtime_context_impl, - writer, pool_, realtime_store_state.value()); + return RealtimePrimaryKeyWriter::Create(schema_, trimmed_primary_keys, key_comparator_, + realtime_store_state.value(), restore_max_seq_number, + writer, pool_); } Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index 45462ea6e..cbd2189fc 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -19,9 +19,13 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include +#include #include #include +#include +#include #include +#include #include #include @@ -44,6 +48,7 @@ #include "paimon/common/utils/path_util.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/operation/restore_files.h" +#include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/file_store_commit.h" @@ -52,7 +57,9 @@ #include "paimon/format/file_format_factory.h" #include "paimon/format/reader_builder.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" @@ -61,6 +68,113 @@ #include "paimon/write_context.h" namespace paimon::test { +namespace { + +class TestingMemoryPool final : public MemoryPool { + public: + void* Malloc(uint64_t size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* pointer, size_t old_size, size_t new_size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Realloc(pointer, old_size, new_size, alignment); + } + + void Free(void* pointer, uint64_t size) override { + delegate_->Free(pointer, size); + } + + void Free(void* pointer, uint64_t size, uint64_t alignment) override { + delegate_->Free(pointer, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + bool reject_allocations = false; + int64_t allocation_count = 0; + + private: + std::unique_ptr delegate_ = GetMemoryPool(); +}; + +class FailOnceRealtimeStore final : public RealtimeStore { + public: + FailOnceRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr& fail_next_write) + : delegate_(delegate), fail_next_write_(fail_next_write) {} + + Status Write(RealtimeWriteBatch&& batch) override { + if (*fail_next_write_) { + *fail_next_write_ = false; + return Status::Invalid("injected real-time store write failure"); + } + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + std::shared_ptr fail_next_write_; +}; + +class FailOnceRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit FailOnceRealtimeStoreFactory(const std::shared_ptr& fail_next_write) + : fail_next_write_(fail_next_write) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate, fail_next_write_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr fail_next_write_; +}; + +} class KeyValueFileStoreWriteTest : public ::testing::Test { protected: @@ -128,14 +242,15 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } - std::unique_ptr MakeBatch(const std::shared_ptr& schema, - const std::string& json) const { + std::unique_ptr MakeBatch( + const std::shared_ptr& schema, const std::string& json, + const std::vector& row_kinds = {}) const { auto struct_type = arrow::struct_(schema->fields()); auto array = arrow::ipc::internal::json::ArrayFromJSON(struct_type, json).ValueOrDie(); ::ArrowArray arrow_array; EXPECT_TRUE(arrow::ExportArray(*array, &arrow_array).ok()); RecordBatchBuilder batch_builder(&arrow_array); - return batch_builder.SetBucket(0).Finish().value(); + return batch_builder.SetRowKinds(row_kinds).SetBucket(0).Finish().value(); } std::vector> WriteAndPrepare( @@ -194,6 +309,58 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { EXPECT_NE(nullptr, metadata); return MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy()).value(); } + + Result>> + ReadPreparedRows(const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + context->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected exactly one real-time store"); + } + RealtimeQueryContext query_context{nullptr, nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + views[0].store->CreateQueryReaders( + views[0].read_view, 0, query_context)); + std::vector> rows; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(array); + if (!values || values->num_fields() != 5) { + return Status::Invalid("unexpected prepared real-time batch"); + } + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(values->field(0)); + std::shared_ptr sequences = + std::dynamic_pointer_cast(values->field(1)); + std::shared_ptr offsets = + std::dynamic_pointer_cast(values->field(2)); + std::shared_ptr ids = + std::dynamic_pointer_cast(values->field(3)); + std::shared_ptr payloads = + std::dynamic_pointer_cast(values->field(4)); + if (!row_kinds || !sequences || !offsets || !ids || !payloads) { + return Status::Invalid("unexpected prepared real-time column type"); + } + for (int64_t row = 0; row < values->length(); ++row) { + rows.emplace_back(row_kinds->Value(row), ids->Value(row), + payloads->GetString(row), sequences->Value(row), + offsets->Value(row)); + } + } + reader->Close(); + } + return rows; + } }; TEST_F(KeyValueFileStoreWriteTest, TestWriteWithInvalidBatch) { @@ -310,7 +477,7 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { {Options::WRITE_BUFFER_SIZE, "1"}, }; const std::shared_ptr schema = arrow::schema({ - arrow::field("id", arrow::int64(), /*nullable=*/false), + arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), }); std::unique_ptr dir = UniqueTestDirectory::Create(); @@ -329,13 +496,23 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, FileStoreWrite::Create(std::move(write_context))); - ASSERT_OK(writer->Write(MakeBatch(schema, R"([ + std::unique_ptr batch = + MakeBatch(schema, R"([ [1, "old"], [2, "two"], [1, "new"] - ])"))); + ])", + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER}); + ASSERT_OK(writer->Write(std::move(batch))); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ( + (std::vector{{0, 1, "old", 0, 0}, {2, 1, "new", 2, 2}, {3, 2, "two", 1, 1}}), + prepared_rows); ASSERT_OK_AND_ASSIGN(std::vector progresses, - writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + writer->PrepareCommitWithProgress(0)); ASSERT_EQ(1, progresses.size()); ASSERT_EQ(OffsetRange(0, 3), progresses[0].offset_range); std::shared_ptr commit_message = @@ -351,6 +528,189 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { ASSERT_OK(writer->Close()); } +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeOffsetCollision) { + const std::map options = {{Options::BUCKET, "1"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("_REALTIME_OFFSET", arrow::int64()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(PathUtil::JoinPath(dir->Str(), "foo.db/bar"), "test"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[1, 10]])")), + "PK real-time write schema contains reserved transport field"); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestWriteFailureKeepsCursors) { + const std::map options = { + {Options::BUCKET, "1"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + }; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + auto fail_next_write = std::make_shared(true); + auto factory = std::make_shared(fail_next_write); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[9, "rejected"]])")), + "injected real-time store write failure"); + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "kept"]])"))); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "kept", 0, 0}}), prepared_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { + const std::map options = {{Options::BUCKET, "1"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + std::shared_ptr pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + const int64_t allocations_before_write = pool->allocation_count; + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "one"]])"))); + ASSERT_GT(pool->allocation_count, allocations_before_write); + ASSERT_OK(writer->Close()); + writer.reset(); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector retained_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "one", 0, 0}}), retained_rows); + + std::shared_ptr rejecting_pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rejecting_context, + RealtimeContext::Create()); + WriteContextBuilder rejecting_builder(table_path, "rejecting"); + rejecting_builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(rejecting_context) + .WithMemoryPool(rejecting_pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_write_context, + rejecting_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_writer, + FileStoreWrite::Create(std::move(rejecting_write_context))); + ASSERT_OK(rejecting_writer->Write(MakeBatch(schema, "[]"))); + const int64_t rejecting_allocations_before_write = rejecting_pool->allocation_count; + rejecting_pool->reject_allocations = true; + ASSERT_NOK_WITH_MSG(rejecting_writer->Write(MakeBatch(schema, R"([[2, "two"]])")), + "Out of memory"); + ASSERT_GT(rejecting_pool->allocation_count, rejecting_allocations_before_write); + ASSERT_OK_AND_ASSIGN(std::vector rejected_rows, + ReadPreparedRows(rejecting_context)); + ASSERT_TRUE(rejected_rows.empty()); + ASSERT_OK(rejecting_writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { + const int64_t max = std::numeric_limits::max(); + const std::map options = {{Options::BUCKET, "1"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr initial_context, + RealtimeContext::Create()); + WriteContextBuilder initial_builder(table_path, "initial"); + initial_builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext( + initial_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_write_context, + initial_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_writer, + FileStoreWrite::Create(std::move(initial_write_context))); + ASSERT_OK(initial_writer->Write(MakeBatch(schema, R"([[0, "initial"]])"))); + ASSERT_OK_AND_ASSIGN(std::vector initial_progress, + initial_writer->PrepareCommitWithProgress(0)); + ASSERT_EQ(1, initial_progress.size()); + std::shared_ptr initial_message = + std::dynamic_pointer_cast(initial_progress[0].commit_message); + ASSERT_NE(nullptr, initial_message); + ASSERT_EQ(1, initial_message->GetNewFilesIncrement().NewFiles().size()); + initial_message->GetNewFilesIncrement().NewFiles()[0]->AssignSequenceNumber(max - 2, max - 2); + initial_progress[0].offset_range = OffsetRange(0, max - 1); + + CommitContextBuilder commit_builder(table_path, "initial"); + commit_builder.SetOptions(options); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, commit_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committer, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + committer->CommitWithProgress(initial_progress, 0, std::nullopt)); + ASSERT_OK(initial_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "boundary"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "legal"]])"))); + using PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[2, "overflow"]])")), + "real-time offset range exceeds INT64_MAX"); + ASSERT_OK_AND_ASSIGN(prepared_rows, ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context_impl, + RealtimeContextImpl::Cast(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::vector views, + context_impl->AcquireReadViews()); + ASSERT_EQ(1, views.size()); + ASSERT_EQ(std::optional(OffsetRange(max - 1, max)), + views[0].read_view->GetOffsetRange()); + ASSERT_OK(writer->Close()); + ASSERT_GE(snapshot_id, 1); +} + TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenDefaultCompactRewriterPathEnabled) { ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index 8d8367e39..2f64f6df8 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -81,8 +81,6 @@ class MergeFunctionWrapper; namespace { -/// Concatenates merge readers whose key ranges are ordered and non-overlapping, preserving one -/// projection pipeline without merging independent disk-only components. class ConcatNonOverlappingMergeReaders final : public SortMergeReader { public: explicit ConcatNonOverlappingMergeReaders( @@ -117,7 +115,7 @@ class ConcatNonOverlappingMergeReaders final : public SortMergeReader { size_t current_ = 0; }; -} // namespace +} class MergeFileSplitRead::RealtimeReaderBuilder { public: @@ -219,8 +217,8 @@ class MergeFileSplitRead::RealtimeReaderBuilder { inputs_.reserve(inputs_.size() + additional_readers.size()); for (AdditionalKeyValueReader& additional : additional_readers) { has_unknown_range_ |= !additional.min_key || !additional.max_key; - inputs_.push_back(RangeInput{additional.min_key, additional.max_key, - /*disk_runs=*/{}, std::move(additional.reader)}); + inputs_.push_back(RangeInput{additional.min_key, additional.max_key, {}, + std::move(additional.reader)}); } } @@ -310,7 +308,7 @@ class MergeFileSplitRead::RealtimeReaderBuilder { component.front().disk_runs, first_split_->Partition(), dv_factory_, component.front().disk_runs.size() == 1 ? owner_->context_->GetPredicate() : owner_->predicate_for_keys_, - data_file_path_factory_, /*drop_delete=*/false)); + data_file_path_factory_, false)); component_readers.push_back(std::move(disk_component)); continue; } diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 8c541ec6f..6c1399978 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -128,7 +128,6 @@ class MergeFileSplitRead : public AbstractSplitRead { return key_schema_; } - /// Merges ordinary disk splits with generic additional sorted KeyValue readers. Result> CreateRealtimeReader( const std::vector>& disk_splits, std::vector&& additional_readers); diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index e6e22edfd..4cfdb4c3d 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -21,14 +21,9 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" -#include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" -#include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/scope_guard.h" -#include "paimon/core/core_options.h" -#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" -#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/realtime/arrow_realtime_store.h" #include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" @@ -55,31 +50,8 @@ Result> ArrowRealtimeStoreFactory::Create( request.memory_pool, arrow_pool); } - const PrimaryKeyRealtimeStoreCreateConfig& primary_key_config = - std::get(request.mode_config); - std::vector key_fields; - key_fields.reserve(primary_key_config.primary_keys.size()); - for (const std::string& primary_key : primary_key_config.primary_keys) { - const int32_t field_index = imported_schema->GetFieldIndex(primary_key); - if (field_index < 0) { - return Status::Invalid("primary key ", primary_key, " is missing from write schema"); - } - key_fields.emplace_back(field_index, imported_schema->field(field_index)); - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, - FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); - auto merge_function_wrapper_factory = []() { - auto merge_function = std::make_unique( - /*ignore_delete=*/false); - return std::make_shared(std::move(merge_function)); - }; - PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(request.options)); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(imported_schema, primary_key_config.primary_keys, - key_comparator, merge_function_wrapper_factory, - primary_key_config.restore_max_sequence_number, - core_options.GetReadBatchSize(), request.memory_pool)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool)); return std::shared_ptr(std::move(store)); } diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 8f51c1b1e..0d6de9f5f 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -9,41 +9,24 @@ * * 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. + * 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/realtime/primary_key_realtime_store.h" -#include -#include #include #include #include "arrow/api.h" #include "arrow/c/bridge.h" -#include "paimon/common/data/binary_row_writer.h" -#include "paimon/common/data/columnar/columnar_row_ref.h" #include "paimon/common/metrics/metrics_impl.h" -#include "paimon/common/table/special_fields.h" -#include "paimon/common/types/data_field.h" -#include "paimon/common/types/row_kind.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/core/core_options.h" -#include "paimon/core/io/key_value_in_memory_record_reader.h" -#include "paimon/core/io/key_value_projection_consumer.h" -#include "paimon/core/io/key_value_projection_reader.h" -#include "paimon/core/io/merged_key_value_record_reader.h" -#include "paimon/core/key_value.h" -#include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h" -#include "paimon/core/utils/nested_projection_utils.h" #include "paimon/macros.h" namespace paimon { @@ -83,562 +66,255 @@ Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { namespace { uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { - uint64_t result = 0; + uint64_t total = 0; for (const std::shared_ptr& buffer : data->buffers) { if (buffer) { - result += static_cast(buffer->size()); + total += static_cast(buffer->size()); } } for (const std::shared_ptr& child : data->child_data) { - result += GetArrayMemoryUsage(child); + total += GetArrayMemoryUsage(child); } if (data->dictionary) { - result += GetArrayMemoryUsage(data->dictionary); + total += GetArrayMemoryUsage(data->dictionary); } - return result; -} - -int32_t FindPkQueryFieldIndex(const std::shared_ptr& write_schema, - const std::shared_ptr& read_field) { - Result read_id = NestedProjectionUtils::GetPaimonFieldId(read_field); - if (read_id.ok()) { - Result> write_field = - NestedProjectionUtils::FindFieldByPaimonId(arrow::struct_(write_schema->fields()), - read_id.value()); - if (write_field.ok()) { - return write_schema->GetFieldIndex(write_field.value()->name()); - } - } - - const int32_t name_index = write_schema->GetFieldIndex(read_field->name()); - if (name_index < 0) { - return -1; - } - Result write_id = - NestedProjectionUtils::GetPaimonFieldId(write_schema->field(name_index)); - if (read_id.ok() && write_id.ok() && read_id.value() != write_id.value()) { - return -1; - } - return name_index; + return total; } struct StoredBatch { std::shared_ptr data; - std::vector row_kinds; OffsetRange offset_range; - int64_t first_sequence_number; uint64_t memory_usage; }; -using BatchGroup = std::vector>; class Segment final : public RealtimeSegmentHandle { public: - Segment(const OffsetRange& offset_range, - std::vector>&& batches) - : offset_range_(offset_range), batches_(std::move(batches)) {} + Segment(const OffsetRange& range, std::vector&& batches) + : range_(range), batches_(std::move(batches)) {} OffsetRange GetOffsetRange() const override { - return offset_range_; + return range_; } - - const std::vector>& Batches() const { + const std::vector& Batches() const { return batches_; } - uint64_t GetMemoryUsage() const { - uint64_t result = 0; - for (const std::shared_ptr& batch : batches_) { - result += batch->memory_usage; - } - return result; - } - private: - OffsetRange offset_range_; - std::vector> batches_; + OffsetRange range_; + std::vector batches_; }; -class PrimaryKeyRealtimeReadView final : public RealtimeReadView { +class ReadView final : public RealtimeReadView { public: - explicit PrimaryKeyRealtimeReadView(std::vector&& groups) - : groups_(std::move(groups)) { - if (!groups_.empty()) { - offset_range_ = OffsetRange(groups_.front().front()->offset_range.begin, - groups_.back().back()->offset_range.end); + explicit ReadView(std::vector>&& segments) + : segments_(std::move(segments)) { + if (!segments_.empty()) { + range_ = OffsetRange(segments_.front()->GetOffsetRange().begin, + segments_.back()->GetOffsetRange().end); } } std::optional GetOffsetRange() const override { - return offset_range_; + return range_; } - - const std::vector& Groups() const { - return groups_; + const std::vector>& Segments() const { + return segments_; } private: - std::vector groups_; - std::optional offset_range_; + std::vector> segments_; + std::optional range_; }; -class CommitBatchReader final : public BatchReader { +class RawBatchReader final : public BatchReader { public: - CommitBatchReader(const std::shared_ptr& segment, - const std::shared_ptr& arrow_pool) - : segment_(segment), arrow_pool_(arrow_pool), metrics_(std::make_shared()) {} + RawBatchReader(std::vector batches) + : batches_(std::move(batches)), metrics_(std::make_shared()) {} Result NextBatch() override { - if (!segment_ || next_batch_ >= static_cast(segment_->Batches().size())) { + if (next_ == batches_.size()) { return MakeEofBatch(); } - const std::shared_ptr& stored = segment_->Batches()[next_batch_++]; - const int64_t row_count = stored->data->length(); - arrow::Int8Builder row_kind_builder(arrow_pool_.get()); - PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Reserve(row_count)); - if (stored->row_kinds.empty()) { - for (int64_t i = 0; i < row_count; ++i) { - row_kind_builder.UnsafeAppend(static_cast(RecordBatch::RowKind::INSERT)); - } - } else { - for (RecordBatch::RowKind row_kind : stored->row_kinds) { - row_kind_builder.UnsafeAppend(static_cast(row_kind)); - } - } - std::shared_ptr row_kind_array; - PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); - arrow::ArrayVector arrays = {std::move(row_kind_array)}; - arrays.insert(arrays.end(), stored->data->fields().begin(), stored->data->fields().end()); - arrow::FieldVector fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; - const arrow::FieldVector& value_fields = stored->data->struct_type()->fields(); - fields.insert(fields.end(), value_fields.begin(), value_fields.end()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr output, - arrow::StructArray::Make(arrays, fields)); - auto c_array = std::make_unique(); - auto c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*output, c_array.get(), c_schema.get())); - return ReadBatch(std::move(c_array), std::move(c_schema)); + const std::shared_ptr& batch = batches_[next_++].data; + auto array = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, array.get(), schema.get())); + return ReadBatch(std::move(array), std::move(schema)); } std::shared_ptr GetReaderMetrics() const override { return metrics_; } - void Close() override { - segment_.reset(); + batches_.clear(); } private: - std::shared_ptr segment_; - std::shared_ptr arrow_pool_; + std::vector batches_; + size_t next_ = 0; std::shared_ptr metrics_; - int32_t next_batch_ = 0; -}; - -class KeyRangeBatchReader final : public BatchReader, public PrimaryKeyRangeProvider { - public: - KeyRangeBatchReader(std::unique_ptr&& reader, - const std::shared_ptr& min_key, - const std::shared_ptr& max_key) - : reader_(std::move(reader)), min_key_(min_key), max_key_(max_key) {} - - Result NextBatch() override { - return reader_->NextBatch(); - } - - std::shared_ptr GetReaderMetrics() const override { - return reader_->GetReaderMetrics(); - } - - void Close() override { - reader_->Close(); - } - - std::shared_ptr GetMinKey() const override { - return min_key_; - } - - std::shared_ptr GetMaxKey() const override { - return max_key_; - } - - private: - std::unique_ptr reader_; - std::shared_ptr min_key_; - std::shared_ptr max_key_; }; } // namespace class PrimaryKeyRealtimeStore::Impl { public: - Impl(const std::shared_ptr& write_schema, std::vector primary_keys, - const std::shared_ptr& key_comparator, - const std::function>()>& - merge_function_wrapper_factory, - int64_t next_sequence_number, int32_t read_batch_size, - const std::shared_ptr& memory_pool) - : write_schema_(write_schema), - primary_keys_(std::move(primary_keys)), - key_comparator_(key_comparator), - merge_function_wrapper_factory_(merge_function_wrapper_factory), - next_sequence_number_(next_sequence_number), - read_batch_size_(read_batch_size), - memory_pool_(memory_pool), - arrow_pool_(GetArrowPool(memory_pool)) {} - - Result> CopyKey(const InternalRow& key) const { - auto result = std::make_shared(static_cast(primary_keys_.size())); - BinaryRowWriter writer(result.get(), /*initial_size=*/128, memory_pool_.get()); - writer.Reset(); - for (int32_t index = 0; index < static_cast(primary_keys_.size()); ++index) { - std::shared_ptr field = - write_schema_->GetFieldByName(primary_keys_[index]); - PAIMON_ASSIGN_OR_RAISE(InternalRow::FieldGetterFunc getter, - InternalRow::CreateFieldGetter(index, field->type(), - /*use_view=*/true)); - PAIMON_ASSIGN_OR_RAISE(BinaryRowWriter::FieldSetterFunc setter, - BinaryRowWriter::CreateFieldSetter(index, field->type())); - setter(getter(key), &writer); - } - writer.Complete(); - return std::static_pointer_cast(result); - } - - Result, std::shared_ptr>> GetKeyRange( - const std::shared_ptr& values) const { - arrow::ArrayVector key_arrays; - key_arrays.reserve(primary_keys_.size()); - for (const std::string& primary_key : primary_keys_) { - std::shared_ptr key_array = values->GetFieldByName(primary_key); - if (!key_array) { - return Status::Invalid("primary key is missing from PK query batch: ", primary_key); - } - key_arrays.push_back(std::move(key_array)); - } - auto context = std::make_shared(key_arrays, memory_pool_); - int64_t min_row = 0; - int64_t max_row = 0; - for (int64_t row = 1; row < values->length(); ++row) { - ColumnarRowRef current(context, row); - ColumnarRowRef min_key(context, min_row); - ColumnarRowRef max_key(context, max_row); - if (key_comparator_->CompareTo(current, min_key) < 0) { - min_row = row; - } - if (key_comparator_->CompareTo(current, max_key) > 0) { - max_row = row; - } - } - ColumnarRowRef min_key(context, min_row); - ColumnarRowRef max_key(context, max_row); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr copied_min, CopyKey(min_key)); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr copied_max, CopyKey(max_key)); - return std::make_pair(std::move(copied_min), std::move(copied_max)); - } + explicit Impl(std::shared_ptr prepared_schema) + : prepared_schema_(std::move(prepared_schema)) {} Status Write(RealtimeWriteBatch&& write_batch) { if (!write_batch.batch || !write_batch.batch->GetData()) { return Status::Invalid("PK real-time write batch is null"); } const int64_t row_count = write_batch.batch->GetData()->length; - if (row_count <= 0 || write_batch.offset_range.begin < 0 || - write_batch.offset_range.Count() != row_count) { + if (write_batch.offset_range.begin < 0 || write_batch.offset_range.Count() != row_count || + row_count <= 0) { return Status::Invalid("PK real-time offset range does not match batch row count"); } - const std::vector& row_kinds = write_batch.batch->GetRowKind(); - if (!row_kinds.empty() && static_cast(row_kinds.size()) != row_count) { - return Status::Invalid("PK real-time row-kind count does not match batch row count"); - } - for (RecordBatch::RowKind row_kind : row_kinds) { - PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, - RowKind::FromByteValue(static_cast(row_kind))); - static_cast(validated); - } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr imported, + std::shared_ptr array, arrow::ImportArray(write_batch.batch->GetData(), - arrow::struct_(write_schema_->fields()))); - if (!imported || imported->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("PK real-time write data is not a StructArray"); + arrow::struct_(prepared_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time prepared batch is not a StructArray"); } - std::shared_ptr values = - checked_pointer_cast(imported); - PAIMON_RETURN_NOT_OK_FROM_ARROW(values->ValidateFull()); - + std::shared_ptr prepared = + checked_pointer_cast(array); + PAIMON_RETURN_NOT_OK_FROM_ARROW(prepared->ValidateFull()); std::lock_guard lock(mutex_); if (last_offset_ && write_batch.offset_range.begin != last_offset_.value()) { return Status::Invalid("PK real-time offset ranges must be contiguous"); } - if (row_count > std::numeric_limits::max() - next_sequence_number_) { - return Status::Invalid("PK sequence range exceeds INT64_MAX"); - } - auto stored = std::make_shared( - StoredBatch{std::move(values), row_kinds, write_batch.offset_range, - next_sequence_number_, GetArrayMemoryUsage(imported->data())}); - building_batches_.push_back(std::move(stored)); - building_memory_usage_ += building_batches_.back()->memory_usage; + building_.push_back( + StoredBatch{prepared, write_batch.offset_range, GetArrayMemoryUsage(prepared->data())}); + building_memory_usage_ += building_.back().memory_usage; last_offset_ = write_batch.offset_range.end; - next_sequence_number_ += row_count; return Status::OK(); } Result>> SealForCommit() { std::lock_guard lock(mutex_); - if (building_batches_.empty()) { + if (building_.empty()) { return std::optional>(); } - const OffsetRange range(building_batches_.front()->offset_range.begin, - building_batches_.back()->offset_range.end); - auto segment = std::make_shared(range, std::move(building_batches_)); - sealed_segments_.push_back(segment); - building_batches_.clear(); + OffsetRange range(building_.front().offset_range.begin, building_.back().offset_range.end); + std::shared_ptr segment = std::make_shared(range, std::move(building_)); + sealed_.push_back(segment); + building_.clear(); building_memory_usage_ = 0; return std::optional>(std::move(segment)); } Result>> CreateCommitReaders( - const std::shared_ptr& segment) { - std::shared_ptr typed = std::dynamic_pointer_cast(segment); - if (!typed) { + const std::shared_ptr& handle) { + std::shared_ptr segment = std::dynamic_pointer_cast(handle); + if (!segment) { return Status::Invalid("segment was not created by the PK real-time store"); } - std::vector> result; - result.push_back(std::make_unique(typed, arrow_pool_)); - return result; + std::vector> readers; + readers.reserve(segment->Batches().size()); + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back(std::make_unique(std::vector{batch})); + } + return readers; } Result> AcquireReadView() { std::lock_guard lock(mutex_); - std::vector groups; - groups.reserve(sealed_segments_.size() + (building_batches_.empty() ? 0 : 1)); - for (const std::shared_ptr& segment : sealed_segments_) { - groups.push_back(segment->Batches()); + std::vector> segments = sealed_; + if (!building_.empty()) { + OffsetRange range(building_.front().offset_range.begin, + building_.back().offset_range.end); + segments.push_back( + std::make_shared(range, std::vector(building_))); } - if (!building_batches_.empty()) { - groups.push_back(building_batches_); - } - return std::shared_ptr(new PrimaryKeyRealtimeReadView(std::move(groups))); + return std::shared_ptr(new ReadView(std::move(segments))); } Result>> CreateQueryReaders( - const std::shared_ptr& view, int64_t lower, - const RealtimeQueryContext& context) { - std::shared_ptr typed = - std::dynamic_pointer_cast(view); + const std::shared_ptr& view, int64_t, const RealtimeQueryContext&) { + std::shared_ptr typed = std::dynamic_pointer_cast(view); if (!typed) { return Status::Invalid("read view was not created by the PK real-time store"); } - if (!context.read_schema || !context.read_schema->release) { - return Status::Invalid("PK real-time query read schema is null"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr requested, - arrow::ImportSchema(context.read_schema)); - arrow::FieldVector output_fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; - arrow::FieldVector aligned_value_fields = write_schema_->fields(); - std::vector projection = {KeyValueProjectionConsumer::kValueKindProjection}; - for (const std::shared_ptr& field : requested->fields()) { - if (field->name() == SpecialFields::ValueKind().Name()) { - continue; - } - output_fields.push_back(field); - if (field->name() == SpecialFields::SequenceNumber().Name()) { - projection.push_back(KeyValueProjectionConsumer::kSequenceNumberProjection); - continue; - } - int32_t index = FindPkQueryFieldIndex(write_schema_, field); - if (index < 0) { - Result field_id = NestedProjectionUtils::GetPaimonFieldId(field); - if (!field_id.ok()) { - return Status::Invalid( - "PK real-time query field is missing from write schema: ", field->name()); - } - std::string internal_name = - "__paimon_pk_realtime_null_" + std::to_string(field_id.value()); - while ( - NestedProjectionUtils::FindFieldByName(aligned_value_fields, internal_name)) { - internal_name.push_back('_'); - } - index = static_cast(aligned_value_fields.size()); - aligned_value_fields.push_back(field->WithName(internal_name)); - } else { - aligned_value_fields[index] = write_schema_->field(index)->WithType(field->type()); - } - projection.push_back(index); + std::vector> readers; + size_t batch_count = 0; + for (const std::shared_ptr& segment : typed->Segments()) { + batch_count += segment->Batches().size(); } - const std::shared_ptr aligned_value_type = - arrow::struct_(aligned_value_fields); - - std::vector> result; - for (const BatchGroup& group : typed->Groups()) { - std::vector> batch_readers; - std::shared_ptr min_key; - std::shared_ptr max_key; - for (const std::shared_ptr& batch : group) { - if (batch->offset_range.end <= lower) { - continue; - } - const int64_t offset = std::max(0, lower - batch->offset_range.begin); - const int64_t length = batch->data->length() - offset; - std::shared_ptr sliced = batch->data->Slice(offset, length); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr aligned, - NestedProjectionUtils::AlignArrayToReadType( - sliced, aligned_value_type, arrow_pool_.get())); - if (!aligned || aligned->type_id() != arrow::Type::STRUCT) { - return Status::Invalid( - "PK real-time query projection did not produce a " - "StructArray"); - } - std::shared_ptr selected = - checked_pointer_cast(aligned); - using KeyRange = - std::pair, std::shared_ptr>; - PAIMON_ASSIGN_OR_RAISE(KeyRange key_range, GetKeyRange(selected)); - if (!min_key || key_comparator_->CompareTo(*key_range.first, *min_key) < 0) { - min_key = key_range.first; - } - if (!max_key || key_comparator_->CompareTo(*key_range.second, *max_key) > 0) { - max_key = key_range.second; - } - std::vector selected_kinds; - if (!batch->row_kinds.empty()) { - selected_kinds.assign(batch->row_kinds.begin() + offset, - batch->row_kinds.end()); - } - std::unique_ptr reader = - std::make_unique( - batch->first_sequence_number + offset, selected, selected_kinds, - primary_keys_, /*user_defined_sequence_fields=*/std::vector(), - /*sequence_fields_ascending=*/true, key_comparator_, memory_pool_); - std::shared_ptr> batch_merge = - merge_function_wrapper_factory_(); - if (!batch_merge) { - return Status::Invalid("merge function wrapper factory returned null"); - } - batch_readers.push_back(std::make_unique( - std::move(reader), key_comparator_, batch_merge)); - } - if (batch_readers.empty()) { - continue; - } - std::shared_ptr> group_merge = - merge_function_wrapper_factory_(); - if (!group_merge) { - return Status::Invalid("merge function wrapper factory returned null"); + readers.reserve(batch_count); + for (const std::shared_ptr& segment : typed->Segments()) { + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back( + std::make_unique(std::vector{batch})); } - auto merged = std::make_unique( - std::move(batch_readers), key_comparator_, - /*user_defined_seq_comparator=*/nullptr, group_merge); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr projected, - KeyValueProjectionReader::Create(std::move(merged), arrow::schema(output_fields), - projection, read_batch_size_, memory_pool_)); - result.push_back( - std::make_unique(std::move(projected), min_key, max_key)); } - return result; + return readers; } - Status AdvanceCommittedOffset(int64_t committed_end_offset) { + Status AdvanceCommittedOffset(int64_t committed_end) { std::lock_guard lock(mutex_); - sealed_segments_.erase( - std::remove_if(sealed_segments_.begin(), sealed_segments_.end(), - [committed_end_offset](const std::shared_ptr& segment) { - return segment->GetOffsetRange().end <= committed_end_offset; - }), - sealed_segments_.end()); + while (!sealed_.empty() && sealed_.front()->GetOffsetRange().end <= committed_end) { + sealed_.erase(sealed_.begin()); + } return Status::OK(); } uint64_t GetMemoryUsage() const { std::lock_guard lock(mutex_); - uint64_t result = building_memory_usage_; - for (const std::shared_ptr& segment : sealed_segments_) { - result += segment->GetMemoryUsage(); + uint64_t total = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_) { + for (const StoredBatch& batch : segment->Batches()) { + total += batch.memory_usage; + } } - return result; + return total; } private: - std::shared_ptr write_schema_; - std::vector primary_keys_; - std::shared_ptr key_comparator_; - std::function>()> - merge_function_wrapper_factory_; - int64_t next_sequence_number_; - int32_t read_batch_size_; - std::shared_ptr memory_pool_; - std::shared_ptr arrow_pool_; + std::shared_ptr prepared_schema_; mutable std::mutex mutex_; - std::vector> building_batches_; - std::vector> sealed_segments_; + std::vector building_; + std::vector> sealed_; uint64_t building_memory_usage_ = 0; std::optional last_offset_; }; -Result> PrimaryKeyRealtimeStore::Create( - const std::shared_ptr& write_schema, - const std::vector& primary_keys, - const std::shared_ptr& key_comparator, - const std::function>()>& - merge_function_wrapper_factory, - int64_t restore_max_sequence_number, int32_t read_batch_size, - const std::shared_ptr& memory_pool) { - if (!write_schema || primary_keys.empty() || !key_comparator || - !merge_function_wrapper_factory || !memory_pool || read_batch_size <= 0) { - return Status::Invalid("PK real-time store requires schema, keys, merge helpers, and pool"); - } - if (restore_max_sequence_number < -1) { - return Status::Invalid("PK restore max sequence number must be at least -1"); - } - if (restore_max_sequence_number == std::numeric_limits::max()) { - return Status::Invalid("PK sequence number has reached INT64_MAX"); - } - for (const std::string& key : primary_keys) { - if (write_schema->GetFieldIndex(key) < 0) { - return Status::Invalid("primary key ", key, " is missing from write schema"); - } - } - auto impl = std::make_unique( - write_schema, primary_keys, key_comparator, merge_function_wrapper_factory, - restore_max_sequence_number + 1, read_batch_size, memory_pool); - return std::shared_ptr(new PrimaryKeyRealtimeStore(std::move(impl))); -} - PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) : impl_(std::move(impl)) {} - PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& prepared_schema, + const std::shared_ptr& memory_pool) { + if (!prepared_schema || !memory_pool) { + return Status::Invalid("PK prepared schema or memory pool is null"); + } + return std::shared_ptr( + new PrimaryKeyRealtimeStore(std::make_unique(prepared_schema))); +} Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { return impl_->Write(std::move(batch)); } - Result>> PrimaryKeyRealtimeStore::SealForCommit() { return impl_->SealForCommit(); } - Result>> PrimaryKeyRealtimeStore::CreateCommitReaders( const std::shared_ptr& segment) { return impl_->CreateCommitReaders(segment); } - Result> PrimaryKeyRealtimeStore::AcquireReadView() { return impl_->AcquireReadView(); } - Result>> PrimaryKeyRealtimeStore::CreateQueryReaders( - const std::shared_ptr& view, int64_t offset_begin, + const std::shared_ptr& view, int64_t offset, const RealtimeQueryContext& context) { - return impl_->CreateQueryReaders(view, offset_begin, context); + return impl_->CreateQueryReaders(view, offset, context); } - -Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_offset) { - return impl_->AdvanceCommittedOffset(committed_offset); +Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t offset) { + return impl_->AdvanceCommittedOffset(offset); } - uint64_t PrimaryKeyRealtimeStore::GetMemoryUsage() const { return impl_->GetMemoryUsage(); } diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h index 017864c04..5e18dd74f 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.h +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -19,11 +19,7 @@ #pragma once -#include -#include #include -#include -#include #include "paimon/realtime/realtime_store.h" @@ -34,34 +30,15 @@ class Schema; namespace paimon { class CoreOptions; -class FieldsComparator; -struct KeyValue; class MemoryPool; -class InternalRow; -template -class MergeFunctionWrapper; Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options); -/// Optional metadata exposed by PK query readers with a known inclusive key range. -class PrimaryKeyRangeProvider { - public: - virtual ~PrimaryKeyRangeProvider() = default; - - virtual std::shared_ptr GetMinKey() const = 0; - virtual std::shared_ptr GetMaxKey() const = 0; -}; - -/// In-memory store for primary-key real-time writes. +/// In-memory store for prepared primary-key real-time batches. class PrimaryKeyRealtimeStore final : public RealtimeStore { public: static Result> Create( - const std::shared_ptr& write_schema, - const std::vector& primary_keys, - const std::shared_ptr& key_comparator, - const std::function>()>& - merge_function_wrapper_factory, - int64_t restore_max_sequence_number, int32_t read_batch_size, + const std::shared_ptr& prepared_schema, const std::shared_ptr& memory_pool); ~PrimaryKeyRealtimeStore() override; diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index 66901a6b1..43831d7be 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -9,23 +9,18 @@ * * 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. + * 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/realtime/primary_key_realtime_store.h" -#include -#include -#include #include #include #include -#include #include #include "arrow/api.h" @@ -33,14 +28,52 @@ #include "arrow/ipc/json_simple.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/core_options.h" -#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" -#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/macros.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +std::shared_ptr PreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false), + arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(PreparedSchema()->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +Result ReadJson(const std::vector>& readers) { + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + batches.push_back(std::move(array)); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches)); + return result->ToString(); +} TEST(PrimaryKeyRealtimeStoreOptionsTest, TestSupportedOptions) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); @@ -66,428 +99,69 @@ TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { } } -class PrimaryKeyRealtimeStoreTest : public testing::Test { - public: - void SetUp() override { - pool_ = std::shared_ptr(GetMemoryPool()); - schema_ = arrow::schema( - {arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); - ASSERT_OK_AND_ASSIGN(store_, CreateStore(schema_, {"id"}, /*restore_max_sequence=*/4)); - } - - Result> CreateStore( - const std::shared_ptr& schema, const std::vector& primary_keys, - int64_t restore_max_sequence) const { - std::vector key_fields; - key_fields.reserve(primary_keys.size()); - for (const std::string& primary_key : primary_keys) { - const int32_t index = schema->GetFieldIndex(primary_key); - key_fields.emplace_back(index, schema->field(index)); - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, - FieldsComparator::Create(key_fields, - /*is_ascending_order=*/true)); - auto merge_factory = []() { - auto merge_function = - std::make_unique(/*ignore_delete=*/false); - return std::make_shared(std::move(merge_function)); - }; - return PrimaryKeyRealtimeStore::Create(schema, primary_keys, key_comparator, merge_factory, - restore_max_sequence, - /*read_batch_size=*/2, pool_); - } - - std::unique_ptr MakeBatch( - const std::string& json, const std::vector& row_kinds = {}, - const std::shared_ptr& schema = nullptr) const { - const std::shared_ptr& batch_schema = schema ? schema : schema_; - std::shared_ptr array = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(batch_schema->fields()), json) - .ValueOrDie(); - ArrowArray c_array; - EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); - RecordBatchBuilder builder(&c_array); - builder.SetRowKinds(row_kinds); - return builder.Finish().value(); - } - - std::unique_ptr MakeReadSchema(const arrow::FieldVector& fields) const { - auto c_schema = std::make_unique(); - EXPECT_TRUE(arrow::ExportSchema(*arrow::schema(fields), c_schema.get()).ok()); - return c_schema; - } - - void AssertReaderOutput(const std::vector>& readers, - const std::shared_ptr& type, - const std::string& json) const { - std::vector> batches; - for (const std::unique_ptr& reader : readers) { - while (true) { - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); - if (BatchReader::IsEofBatch(batch)) { - break; - } - arrow::Result> imported = - arrow::ImportArray(batch.first.get(), batch.second.get()); - ASSERT_TRUE(imported.ok()) << imported.status().ToString(); - batches.push_back(std::move(imported).ValueOrDie()); - } - } - ASSERT_FALSE(batches.empty()); - arrow::Result> concatenated = arrow::Concatenate(batches); - ASSERT_TRUE(concatenated.ok()) << concatenated.status().ToString(); - std::shared_ptr actual = std::move(concatenated).ValueOrDie(); - std::shared_ptr expected = - arrow::ipc::internal::json::ArrayFromJSON(type, json).ValueOrDie(); - ASSERT_TRUE(actual->Equals(*expected)) - << "expected: " << expected->ToString() << ", actual: " << actual->ToString(); - for (const std::unique_ptr& reader : readers) { - reader->Close(); - } - } - - std::shared_ptr CommitType() const { - return arrow::struct_({ - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), - schema_->field(0), - schema_->field(1), - }); - } - - std::shared_ptr QueryType() const { - return arrow::struct_({ - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()), - schema_->field(0), - schema_->field(1), - }); - } - - arrow::FieldVector FullQueryFields( - const std::shared_ptr& schema = nullptr) const { - const std::shared_ptr& query_schema = schema ? schema : schema_; - arrow::FieldVector fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; - fields.insert(fields.end(), query_schema->fields().begin(), query_schema->fields().end()); - return fields; - } - - protected: - std::shared_ptr pool_; - std::shared_ptr schema_; - std::shared_ptr store_; -}; - -TEST_F(PrimaryKeyRealtimeStoreTest, TestWriteAndSeal) { +TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::optional> segment, - store_->SealForCommit()); + store->SealForCommit()); ASSERT_FALSE(segment.has_value()); - ASSERT_NOK_WITH_MSG(store_->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), + ASSERT_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), "write batch is null"); ASSERT_NOK_WITH_MSG( - store_->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "a"], [2, "b"]])"), OffsetRange(0, 0)}), + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 1, "one"]])"), OffsetRange(0, 0)}), "offset range does not match batch row count"); - ASSERT_OK( - store_->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "a"], [2, "b"]])"), OffsetRange(0, 2)})); - ASSERT_NOK_WITH_MSG( - store_->Write(RealtimeWriteBatch{MakeBatch(R"([[4, "d"]])"), OffsetRange(3, 4)}), - "offset ranges must be contiguous"); - ASSERT_OK(store_->Write(RealtimeWriteBatch{MakeBatch(R"([[3, "c"]])"), OffsetRange(2, 3)})); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[0, 1, 0, 1, "one"], [0, 2, 1, 2, "two"]])"), OffsetRange(0, 2)})); + ASSERT_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 3, 3, 3, "three"]])"), + OffsetRange(3, 4)}), + "offset ranges must be contiguous"); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 3, 2, 3, "three"]])"), OffsetRange(2, 3)})); - ASSERT_OK_AND_ASSIGN(segment, store_->SealForCommit()); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); ASSERT_TRUE(segment.has_value()); ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); - ASSERT_GT(store_->GetMemoryUsage(), 0); - - struct ValidationCase { - int64_t restore_max_sequence; - std::string error; - }; - const std::vector cases = { - {-2, "restore max sequence number must be at least -1"}, - {std::numeric_limits::max(), "sequence number has reached INT64_MAX"}, - }; - for (const ValidationCase& test_case : cases) { - ASSERT_NOK_WITH_MSG(CreateStore(schema_, {"id"}, test_case.restore_max_sequence), - test_case.error); - } -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { - ASSERT_OK(store_->Write(RealtimeWriteBatch{ - MakeBatch(R"([[3, "three"], [1, "before"]])", - {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_BEFORE}), - OffsetRange(0, 2)})); - ASSERT_OK(store_->Write(RealtimeWriteBatch{ - MakeBatch(R"([[2, "after"]])", {RecordBatch::RowKind::UPDATE_AFTER}), OffsetRange(2, 3)})); - ASSERT_OK(store_->Write( - RealtimeWriteBatch{MakeBatch(R"([[4, "deleted"], [0, "zero"]])", - {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}), - OffsetRange(3, 5)})); - ASSERT_OK_AND_ASSIGN(std::optional> segment, - store_->SealForCommit()); - ASSERT_TRUE(segment.has_value()); - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store_->CreateCommitReaders(segment.value())); - AssertReaderOutput(readers, CommitType(), - R"([[0, 3, "three"], [1, 1, "before"], [2, 2, "after"], - [3, 4, "deleted"], [0, 0, "zero"]])"); - - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - AssertReaderOutput(readers, QueryType(), - R"([[0, 9, 0, "zero"], [1, 6, 1, "before"], [2, 7, 2, "after"], - [0, 5, 3, "three"], [3, 8, 4, "deleted"]])"); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestMutationMerge) { - ASSERT_OK(store_->Write( - RealtimeWriteBatch{MakeBatch(R"([[1, "old"], [2, "two"]])"), OffsetRange(0, 2)})); - ASSERT_OK(store_->Write(RealtimeWriteBatch{ - MakeBatch(R"([[1, "new"], [2, "gone"]])", - {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE}), - OffsetRange(2, 4)})); - - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - AssertReaderOutput(readers, QueryType(), R"([[2, 7, 1, "new"], [3, 8, 2, "gone"]])"); + ASSERT_GT(store->GetMemoryUsage(), 0); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 4, 4, 4, "four"]])"), OffsetRange(4, 5)}), + "offset ranges must be contiguous"); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); } -TEST_F(PrimaryKeyRealtimeStoreTest, TestReadViewLifecycle) { - ASSERT_OK(store_->Write(RealtimeWriteBatch{MakeBatch(R"([[10, "a"], [11, "b"], [12, "c"]])"), - OffsetRange(10, 13)})); +TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[0, 5, 0, 3, "three"], [1, 6, 1, 1, "before"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, - store_->SealForCommit()); + store->SealForCommit()); ASSERT_TRUE(segment.has_value()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - ASSERT_EQ(std::optional(OffsetRange(10, 13)), view->GetOffsetRange()); - - ASSERT_OK(store_->AdvanceCommittedOffset(13)); - ASSERT_EQ(0, store_->GetMemoryUsage()); - ASSERT_OK( - store_->Write(RealtimeWriteBatch{MakeBatch(R"([[13, "later"]])"), OffsetRange(13, 14)})); - - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store_->CreateQueryReaders(view, /*offset_begin=*/11, context)); - AssertReaderOutput(readers, QueryType(), R"([[0, 6, 11, "b"], [0, 7, 12, "c"]])"); - - std::unique_ptr empty_schema = MakeReadSchema(FullQueryFields()); - context.read_schema = empty_schema.get(); - ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/13, context)); - ASSERT_TRUE(readers.empty()); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryKeyRange) { - ASSERT_OK(store_->Write( - RealtimeWriteBatch{MakeBatch(R"([[5, "five"], [1, "one"]])"), OffsetRange(0, 2)})); - ASSERT_OK(store_->SealForCommit()); - ASSERT_OK(store_->Write( - RealtimeWriteBatch{MakeBatch(R"([[9, "nine"], [7, "seven"]])"), OffsetRange(2, 4)})); - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; ASSERT_OK_AND_ASSIGN(std::vector> readers, - store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - ASSERT_EQ(2, readers.size()); - const std::vector> key_ranges = {{1, 5}, {7, 9}}; - for (size_t i = 0; i < readers.size(); ++i) { - auto* range = dynamic_cast(readers[i].get()); - ASSERT_NE(nullptr, range); - ASSERT_EQ(key_ranges[i].first, range->GetMinKey()->GetLong(0)); - ASSERT_EQ(key_ranges[i].second, range->GetMaxKey()->GetLong(0)); - } - AssertReaderOutput(readers, QueryType(), - R"([[0, 6, 1, "one"], [0, 5, 5, "five"], [0, 8, 7, "seven"], - [0, 7, 9, "nine"]])"); - - ASSERT_OK(store_->AdvanceCommittedOffset(2)); - ASSERT_OK_AND_ASSIGN(view, store_->AcquireReadView()); - read_schema = MakeReadSchema(FullQueryFields()); - context.read_schema = read_schema.get(); - ASSERT_OK_AND_ASSIGN(readers, store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - ASSERT_EQ(1, readers.size()); - auto* range = dynamic_cast(readers[0].get()); - ASSERT_NE(nullptr, range); - ASSERT_EQ(7, range->GetMinKey()->GetLong(0)); - ASSERT_EQ(9, range->GetMaxKey()->GetLong(0)); - AssertReaderOutput(readers, QueryType(), R"([[0, 8, 7, "seven"], [0, 7, 9, "nine"]])"); + store->CreateCommitReaders(segment.value())); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_EQ( + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 0,\n 1,\n 2\n ]\n-- " + "child 1 type: int64\n [\n 5,\n 6,\n 7\n ]\n-- child 2 type: int64\n [\n " + "0,\n 1,\n 2\n ]\n-- child 3 type: int64\n [\n 3,\n 1,\n 2\n ]\n-- child " + "4 type: string\n [\n \"three\",\n \"before\",\n \"after\"\n ]", + actual); } -TEST_F(PrimaryKeyRealtimeStoreTest, TestSequenceExhaustion) { - const int64_t max_sequence = std::numeric_limits::max(); +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - CreateStore(schema_, {"id"}, max_sequence - 3)); - ASSERT_OK(store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, "kept"]])"), OffsetRange(10, 11)})); - ASSERT_NOK_WITH_MSG( - store->Write(RealtimeWriteBatch{ - MakeBatch(R"([[7, "rejected-a"], [8, "rejected-b"], [9, "rejected-c"]])"), - OffsetRange(11, 14)}), - "sequence range exceeds INT64_MAX"); + PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); ASSERT_OK( - store->Write(RealtimeWriteBatch{MakeBatch(R"([[2, "also-kept"]])"), OffsetRange(11, 12)})); - + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, store->SealForCommit()); - ASSERT_TRUE(segment.has_value()); - ASSERT_EQ(OffsetRange(10, 12), segment.value()->GetOffsetRange()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); - ASSERT_EQ(std::optional(OffsetRange(10, 12)), view->GetOffsetRange()); - std::unique_ptr read_schema = MakeReadSchema(FullQueryFields()); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store->CreateQueryReaders(view, /*offset_begin=*/10, context)); - AssertReaderOutput(readers, QueryType(), - R"([[0, 9223372036854775805, 1, "kept"], - [0, 9223372036854775806, 2, "also-kept"]])"); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjection) { - ASSERT_OK( - store_->Write(RealtimeWriteBatch{MakeBatch(R"([[2, "b"], [1, "a"]])"), OffsetRange(0, 2)})); - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); - const std::shared_ptr value_kind = - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()); - const std::shared_ptr sequence = - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()); - struct ProjectionCase { - arrow::FieldVector requested; - std::shared_ptr expected_type; - std::string expected_json; - }; - const std::vector cases = { - {{schema_->field(1), value_kind, sequence, schema_->field(0)}, - arrow::struct_({value_kind, schema_->field(1), sequence, schema_->field(0)}), - R"([[0, "a", 6, 1], [0, "b", 5, 2]])"}, - {{schema_->field(0), value_kind}, - arrow::struct_({value_kind, schema_->field(0)}), - R"([[0, 1], [0, 2]])"}, - }; - for (const ProjectionCase& test_case : cases) { - std::unique_ptr read_schema = MakeReadSchema(test_case.requested); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store_->CreateQueryReaders(view, /*offset_begin=*/0, context)); - AssertReaderOutput(readers, test_case.expected_type, test_case.expected_json); - } - - std::unique_ptr read_schema = - MakeReadSchema({arrow::field("unknown", arrow::int64())}); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_NOK_WITH_MSG(store_->CreateQueryReaders(view, /*offset_begin=*/0, context), - "query field is missing from write schema: unknown"); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestQueryProjectionMatchesRenamedFieldsById) { - const std::shared_ptr id = - DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))); - const std::shared_ptr value = - DataField::ConvertDataFieldToArrowField(DataField(1, arrow::field("value", arrow::utf8()))); - const std::shared_ptr write_schema = arrow::schema({id, value}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - CreateStore(write_schema, {"id"}, /*restore_max_sequence=*/4)); - ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[1, "kept"]])", {}, write_schema), OffsetRange(0, 1)})); - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); - - const std::shared_ptr renamed_value = DataField::ConvertDataFieldToArrowField( - DataField(1, arrow::field("renamed", arrow::utf8()))); - const std::shared_ptr renamed_id = DataField::ConvertDataFieldToArrowField( - DataField(0, arrow::field("renamed_id", arrow::int64()))); - const std::shared_ptr replaced = - DataField::ConvertDataFieldToArrowField(DataField(2, arrow::field("value", arrow::utf8()))); - const std::shared_ptr replaced_id = - DataField::ConvertDataFieldToArrowField(DataField(4, arrow::field("id", arrow::int64()))); - const std::shared_ptr added = - DataField::ConvertDataFieldToArrowField(DataField(3, arrow::field("added", arrow::utf8()))); - std::unique_ptr read_schema = - MakeReadSchema({renamed_value, renamed_id, replaced, replaced_id, added}); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store->CreateQueryReaders(view, /*offset_begin=*/0, context)); - const std::shared_ptr result_type = - arrow::struct_({DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), - renamed_value, renamed_id, replaced, replaced_id, added}); - AssertReaderOutput(readers, result_type, R"([[0, "kept", 1, null, null, null]])"); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestNestedProjection) { - const std::shared_ptr id = - DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))); - const std::shared_ptr a = - DataField::ConvertDataFieldToArrowField(DataField(10, arrow::field("a", arrow::int64()))); - const std::shared_ptr b = - DataField::ConvertDataFieldToArrowField(DataField(11, arrow::field("b", arrow::int64()))); - const std::shared_ptr payload = DataField::ConvertDataFieldToArrowField( - DataField(1, arrow::field("payload", arrow::struct_({a, b})))); - const std::shared_ptr nested_schema = arrow::schema({id, payload}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - CreateStore(nested_schema, {"id"}, /*restore_max_sequence=*/4)); - ASSERT_OK(store->Write(RealtimeWriteBatch{ - MakeBatch(R"([[2, [200, 2000]], [1, [100, null]], [3, [300, 3000]]])", {}, nested_schema), - OffsetRange(0, 3)})); - - ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); - const std::shared_ptr projected_payload = payload->WithType(arrow::struct_({b})); - std::unique_ptr read_schema = MakeReadSchema({projected_payload}); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store->CreateQueryReaders(view, /*offset_begin=*/0, context)); - const std::shared_ptr result_type = arrow::struct_( - {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), projected_payload}); - AssertReaderOutput(readers, result_type, R"([[0, [null]], [0, [2000]], [0, [3000]]])"); -} - -TEST_F(PrimaryKeyRealtimeStoreTest, TestCompositeKeyClipping) { - std::shared_ptr composite_schema = - arrow::schema({arrow::field("id", arrow::int64()), arrow::field("region", arrow::utf8()), - arrow::field("value", arrow::utf8())}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - CreateStore(composite_schema, {"id", "region"}, - /*restore_max_sequence=*/4)); - ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[9, "z", "clipped"], [2, "b", "two-b"], [1, "c", "one-c"], - [2, "a", "two-a"]])", - {}, composite_schema), - OffsetRange(20, 24)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); - const std::shared_ptr sequence = - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()); - std::unique_ptr read_schema = - MakeReadSchema({sequence, composite_schema->field(0), composite_schema->field(2)}); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - ASSERT_OK_AND_ASSIGN(std::vector> readers, - store->CreateQueryReaders(view, /*offset_begin=*/21, context)); - ASSERT_EQ(1, readers.size()); - auto* range = dynamic_cast(readers[0].get()); - ASSERT_NE(nullptr, range); - ASSERT_EQ(1, range->GetMinKey()->GetLong(0)); - ASSERT_EQ("c", range->GetMinKey()->GetString(1).ToString()); - ASSERT_EQ(2, range->GetMaxKey()->GetLong(0)); - ASSERT_EQ("b", range->GetMaxKey()->GetString(1).ToString()); - std::shared_ptr query_type = - arrow::struct_({DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()), - sequence, composite_schema->field(0), composite_schema->field(2)}); - AssertReaderOutput(readers, query_type, - R"([[0, 7, 1, "one-c"], [0, 8, 2, "two-a"], - [0, 6, 2, "two-b"]])"); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_EQ(std::optional(OffsetRange(4, 5)), view->GetOffsetRange()); } +} // namespace } // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 066e54e8a..b73cfdb8a 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -42,7 +42,6 @@ #include "paimon/status.h" namespace paimon { - RealtimeContextImpl::RealtimeContextImpl(const std::shared_ptr& factory) : factory_(factory) {} @@ -83,26 +82,6 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( std::lock_guard registry_lock(mutex_); const RealtimePartitionBucket key(request.partition, request.bucket); auto iter = stores_.find(key); - std::optional initial_max_sequence_number; - PrimaryKeyRealtimeStoreCreateConfig* primary_key_config = - std::get_if(&request.mode_config); - if (primary_key_config) { - auto [sequence_iter, inserted] = materialized_max_sequence_numbers_.emplace( - key, primary_key_config->restore_max_sequence_number); - if (!inserted && primary_key_config->restore_max_sequence_number > sequence_iter->second) { - if (iter != stores_.end()) { - if (request.write_schema) { - ArrowSchemaRelease(request.write_schema.get()); - } - return Status::Invalid( - "restore max sequence number exceeds the materialized watermark of an " - "existing PK real-time store"); - } - sequence_iter->second = primary_key_config->restore_max_sequence_number; - } - initial_max_sequence_number = sequence_iter->second; - primary_key_config->restore_max_sequence_number = sequence_iter->second; - } int64_t initial_offset = 0; auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { @@ -134,7 +113,13 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset, initial_max_sequence_number}; + return RealtimeStoreState{iter->second, initial_offset}; + } + if (!request.memory_pool) { + if (request.write_schema) { + ArrowSchemaRelease(request.write_schema.get()); + } + return Status::Invalid("real-time store memory pool is null"); } Result> store_result = factory_->Create(std::move(request)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); @@ -142,17 +127,7 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); } - return RealtimeStoreState{std::move(store), initial_offset, initial_max_sequence_number}; -} - -void RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( - const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { - std::lock_guard lock(mutex_); - auto [iter, inserted] = - materialized_max_sequence_numbers_.emplace(partition_bucket, max_sequence_number); - if (!inserted && max_sequence_number > iter->second) { - iter->second = max_sequence_number; - } + return RealtimeStoreState{std::move(store), initial_offset}; } Result> RealtimeContextImpl::AcquireReadViews() { diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index f4cd3866e..4f62cf1ee 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -47,7 +47,6 @@ class MemoryPool; struct RealtimeStoreState { std::shared_ptr store; int64_t initial_offset; - std::optional initial_max_sequence_number; }; struct RealtimePartitionBucketView { @@ -68,9 +67,6 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); - void AdvanceMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, - int64_t max_sequence_number); - Result> AcquireReadViews(); Result PinReadView(const RealtimePartitionBucketView& view, int64_t ttl_millis); @@ -100,7 +96,6 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::mutex mutex_; std::mutex progress_mutex_; std::map> stores_; - std::map materialized_max_sequence_numbers_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; // Progress already reflected in stores owned by this context. diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index ab0abe4a7..07bbf555b 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -9,12 +9,11 @@ * * 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. + * 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 @@ -31,7 +30,6 @@ #include "arrow/c/helpers.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/memory/memory_pool.h" -#include "paimon/realtime/realtime_store.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -49,26 +47,21 @@ class TestingRealtimeStore : public RealtimeStore { Status Write(RealtimeWriteBatch&&) override { return Status::OK(); } - Result>> SealForCommit() override { return std::optional>(); } - Result>> CreateCommitReaders( const std::shared_ptr&) override { return std::vector>(); } - Result> AcquireReadView() override { ++acquire_count; return std::make_shared(); } - Result>> CreateQueryReaders( const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { return std::vector>(); } - Status AdvanceCommittedOffset(int64_t committed_offset) override { ++advance_count; if (fail_next_advance) { @@ -78,7 +71,6 @@ class TestingRealtimeStore : public RealtimeStore { committed_offsets.push_back(committed_offset); return Status::OK(); } - uint64_t GetMemoryUsage() const override { return 0; } @@ -105,11 +97,11 @@ class TestingRealtimeStoreFactory : public RealtimeStoreFactory { }; std::unique_ptr MakeWriteSchema() { - auto c_schema = std::make_unique(); + auto schema = std::make_unique(); EXPECT_TRUE( - arrow::ExportSchema(*arrow::schema({arrow::field("id", arrow::int64())}), c_schema.get()) + arrow::ExportSchema(*arrow::schema({arrow::field("id", arrow::int64())}), schema.get()) .ok()); - return c_schema; + return schema; } Result> CreateContext( @@ -129,41 +121,29 @@ Result GetOrCreateAppendStore( AppendRealtimeStoreCreateConfig{StatisticsMode::NONE}}); } -Result GetOrCreatePrimaryKeyStore( - const std::shared_ptr& context, - const std::map& partition, int32_t bucket, - int64_t restore_max_sequence_number, const std::shared_ptr& memory_pool) { - return context->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ - MakeWriteSchema(), /*options=*/{}, memory_pool, partition, bucket, - PrimaryKeyRealtimeStoreCreateConfig{{"id"}, restore_max_sequence_number}}); -} - -TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { +TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - - ASSERT_OK_AND_ASSIGN(RealtimeStoreState first_state, + ASSERT_OK_AND_ASSIGN(RealtimeStoreState first, GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, - MakeWriteSchema(), {{"k", "v"}}, pool)); - ASSERT_EQ(0, first_state.initial_offset); - ASSERT_FALSE(first_state.initial_max_sequence_number); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_again_state, - GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, pool)); - ASSERT_EQ(first_state.store, first_again_state.store); - ASSERT_EQ(0, first_again_state.initial_offset); + MakeWriteSchema(), {{"k", "v"}}, GetDefaultPool())); + ASSERT_EQ(0, first.initial_offset); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState second, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_EQ(first.store, second.store); + ASSERT_EQ(0, second.initial_offset); ASSERT_EQ(1, factory->stores.size()); ASSERT_EQ(1, factory->stores[0]->acquire_count); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState second_state, - GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 1, MakeWriteSchema(), {}, pool)); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState third_state, - GetOrCreateAppendStore(context, {{"dt", "2026-08-03"}}, 0, MakeWriteSchema(), {}, pool)); - ASSERT_NE(first_state.store, second_state.store); - ASSERT_NE(first_state.store, third_state.store); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState third, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 1, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState fourth, + GetOrCreateAppendStore(context, {{"dt", "2026-08-03"}}, 0, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_NE(first.store, third.store); + ASSERT_NE(first.store, fourth.store); ASSERT_EQ(3, factory->stores.size()); ASSERT_OK_AND_ASSIGN(std::vector views, @@ -171,57 +151,22 @@ TEST(RealtimeContextTest, TestReusesIndexerAndCapturesRegisteredViews) { ASSERT_EQ(3, views.size()); const RealtimePartitionBucket expected_partition_bucket({{"dt", "2026-08-02"}}, 0); ASSERT_EQ(expected_partition_bucket, views[0].partition_bucket); - ASSERT_EQ(first_state.store, views[0].store); + ASSERT_EQ(first.store, views[0].store); ASSERT_TRUE(views[0].read_view); ASSERT_EQ(2, factory->stores[0]->acquire_count); ASSERT_EQ(1, factory->stores[1]->acquire_count); ASSERT_EQ(1, factory->stores[2]->acquire_count); } -TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { - auto factory = std::make_shared(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - const std::map partition = {{"dt", "2026-08-02"}}; - - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_state, - GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, - /*restore_max_sequence_number=*/4, GetDefaultPool())); - ASSERT_EQ(4, first_state.initial_max_sequence_number); - - const RealtimePartitionBucket partition_bucket(partition, /*bucket=*/0); - context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, /*max_sequence_number=*/8); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState retained_state, - GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, - /*restore_max_sequence_number=*/6, GetDefaultPool())); - ASSERT_EQ(first_state.store, retained_state.store); - ASSERT_EQ(8, retained_state.initial_max_sequence_number); - - ASSERT_NOK_WITH_MSG( - GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/0, - /*restore_max_sequence_number=*/10, GetDefaultPool()), - "restore max sequence number exceeds the materialized watermark of an " - "existing PK real-time store"); - - const RealtimePartitionBucket new_partition_bucket(partition, /*bucket=*/1); - context->AdvanceMaterializedMaxSequenceNumber(new_partition_bucket, - /*max_sequence_number=*/8); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState new_state, - GetOrCreatePrimaryKeyStore(context, partition, /*bucket=*/1, - /*restore_max_sequence_number=*/10, GetDefaultPool())); - ASSERT_EQ(10, new_state.initial_max_sequence_number); -} - TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, pool)); - ASSERT_OK(GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(2, factory->stores.size()); ASSERT_NOK_WITH_MSG(context->AdvanceCommittedProgress(-1, {}), @@ -238,9 +183,9 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); ASSERT_TRUE(factory->stores[1]->committed_offsets.empty()); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState restored_state, - GetOrCreateAppendStore(context, {{"dt", "unknown"}}, 0, MakeWriteSchema(), {}, pool)); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState restored_state, + GetOrCreateAppendStore(context, {{"dt", "unknown"}}, 0, MakeWriteSchema(), + {}, GetDefaultPool())); ASSERT_EQ(9, restored_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress( @@ -260,12 +205,14 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, pool)); - ASSERT_OK(GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); - ASSERT_OK(GetOrCreateAppendStore(context, partition, 2, MakeWriteSchema(), {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 2, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(3, factory->stores.size()); factory->stores[1]->fail_next_advance = true; @@ -281,7 +228,7 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_OK_AND_ASSIGN( RealtimeStoreState failed_store_state, - GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, pool)); + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(8, failed_store_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress(5, committed_offsets)); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 65bcebcad..c85ff6322 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -9,12 +9,11 @@ * * 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. + * 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/realtime/realtime_primary_key_writer.h" @@ -26,95 +25,254 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" -#include "paimon/common/reader/concat_batch_reader.h" +#include "arrow/compute/api.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/types/row_kind.h" -#include "paimon/common/utils/arrow/arrow_utils.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/scope_guard.h" -#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/io/prepared_key_value_reader.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/mergetree/merge_tree_writer.h" #include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_fields.h" #include "paimon/core/utils/commit_increment.h" #include "paimon/macros.h" -#include "paimon/realtime/realtime_context.h" namespace paimon { +namespace { + +struct PreparedArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleasePreparedArray(ArrowArray* array) { + auto* data = static_cast(array->private_data); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); + delete data; +} + +Status RetainPreparedArrayPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release || !arrow_pool) { + return Status::Invalid("cannot retain prepared batch memory pool"); + } + array->private_data = + new PreparedArrayPrivateData{array->release, array->private_data, arrow_pool}; + array->release = ReleasePreparedArray; + return Status::OK(); +} + +Result> PrepareBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, int64_t first_sequence_number, + int64_t first_offset, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr input, + arrow::ImportArray(batch->GetData(), arrow::struct_(write_schema->fields()))); + if (!input || input->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr values = checked_pointer_cast(input); + const int64_t count = values->length(); + arrow::Int8Builder kinds(arrow_pool); + arrow::Int64Builder sequences(arrow_pool); + arrow::Int64Builder offsets(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count)); + const std::vector& row_kinds = batch->GetRowKind(); + for (int64_t row = 0; row < count; ++row) { + const RecordBatch::RowKind kind = + row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row]; + kinds.UnsafeAppend(static_cast(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr kind_array; + std::shared_ptr sequence_array; + std::shared_ptr offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array)); + arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array), + std::move(offset_array)}; + columns.insert(columns.end(), values->fields().begin(), values->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr prepared, + arrow::StructArray::Make(std::move(columns), prepared_schema->fields())); + + std::vector sort_keys; + sort_keys.reserve(trimmed_primary_keys.size() + 1); + for (const std::string& key : trimmed_primary_keys) { + sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + } + sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), + arrow::compute::SortOrder::Ascending); + arrow::compute::ExecContext context(arrow_pool); + arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum indices, + arrow::compute::SortIndices(arrow::Datum(prepared), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(prepared), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + std::shared_ptr sorted_array = sorted.make_array(); + if (!sorted_array || sorted_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time sorted batch is not a StructArray"); + } + return checked_pointer_cast(std::move(sorted_array)); +} + +} // namespace + Result> RealtimePrimaryKeyWriter::Create( - const std::map& partition, int32_t bucket, const std::shared_ptr& write_schema, - const std::shared_ptr& realtime_context, - const std::shared_ptr& merge_tree_writer, - const std::shared_ptr& memory_pool, const RealtimeStoreState& store_state) { - return std::shared_ptr( - new RealtimePrimaryKeyWriter(store_state.store, merge_tree_writer, realtime_context, - RealtimePartitionBucket(partition, bucket), write_schema, - store_state.initial_offset, memory_pool)); + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const RealtimeStoreState& store_state, + int64_t restored_max_sequence_number, const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool) { + if (!store_state.store || !merge_tree_writer || !write_schema || !key_comparator || + !memory_pool) { + return Status::Invalid("PK real-time writer received a null dependency"); + } + if (trimmed_primary_keys.empty()) { + return Status::Invalid("PK real-time writer requires at least one primary key"); + } + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr field = write_schema->GetFieldByName(key); + if (!field) { + return Status::Invalid("PK field is missing from write schema: ", key); + } + key_fields.push_back(std::move(field)); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), write_schema->fields().begin(), + write_schema->fields().end()); + return std::shared_ptr(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, write_schema, + arrow::schema(std::move(prepared_fields)), arrow::schema(std::move(key_fields)), + trimmed_primary_keys, key_comparator, store_state.initial_offset, + restored_max_sequence_number, memory_pool)); } RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( const std::shared_ptr& realtime_store, const std::shared_ptr& merge_tree_writer, - const std::shared_ptr& realtime_context, - const RealtimePartitionBucket& partition_bucket, - const std::shared_ptr& write_schema, int64_t next_offset, - const std::shared_ptr& memory_pool) + const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, int64_t next_offset, + int64_t last_sequence_number, const std::shared_ptr& memory_pool) : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(realtime_store), merge_tree_writer_(merge_tree_writer), - realtime_context_(realtime_context), - partition_bucket_(partition_bucket), write_schema_(write_schema), - next_offset_(next_offset) {} + prepared_schema_(prepared_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { if (!batch || !batch->GetData()) { return Status::Invalid("PK real-time write batch is null"); } - const int64_t row_count = batch->GetData()->length; - if (row_count == 0) { + const int64_t count = batch->GetData()->length; + if (count == 0) { return Status::OK(); } + const std::vector& row_kinds = batch->GetRowKind(); + if (!row_kinds.empty() && static_cast(row_kinds.size()) != count) { + return Status::Invalid("PK real-time row-kind count does not match batch row count"); + } + for (RecordBatch::RowKind row_kind : row_kinds) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, + RowKind::FromByteValue(static_cast(row_kind))); + static_cast(validated); + } std::lock_guard lock(realtime_store_mutex_); - if (row_count > std::numeric_limits::max() - next_offset_) { + if (count > std::numeric_limits::max() - next_offset_) { return Status::Invalid("real-time offset range exceeds INT64_MAX"); } - const OffsetRange range(next_offset_, next_offset_ + row_count); - PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{std::move(batch), range})); - next_offset_ += row_count; + // Reserve INT64_MAX as the exhausted sequence-number sentinel. + if (last_sequence_number_ >= std::numeric_limits::max() - count) { + return Status::Invalid("PK sequence range exceeds INT64_MAX"); + } + const int64_t first_sequence = last_sequence_number_ + 1; + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr prepared, + PrepareBatch(std::move(batch), write_schema_, prepared_schema_, trimmed_primary_keys_, + first_sequence, next_offset_, arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*prepared, output.get())); + PAIMON_RETURN_NOT_OK(RetainPreparedArrayPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ + count)})); + next_offset_ += count; + last_sequence_number_ += count; return Status::OK(); } Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { - std::lock_guard lock(prepare_mutex_); + std::lock_guard prepare_lock(prepare_mutex_); std::optional> segment; { - std::lock_guard realtime_store_lock(realtime_store_mutex_); - PAIMON_ASSIGN_OR_RAISE(std::optional> sealed_segment, + std::lock_guard store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> sealed, realtime_store_->SealForCommit()); - segment = std::move(sealed_segment); + segment = std::move(sealed); + } + if (segment && !segment.value()) { + return Status::Invalid("PK real-time store sealed a null segment"); } + std::optional sealed_range; + int64_t expected_raw_row_count = 0; if (segment) { - PAIMON_RETURN_NOT_OK(FlushSegment(segment.value())); + sealed_range = segment.value()->GetOffsetRange(); + if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin || + __builtin_sub_overflow(sealed_range->end, sealed_range->begin, + &expected_raw_row_count)) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), expected_raw_row_count)); } PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, merge_tree_writer_->PrepareCommit(wait_compaction)); if (segment) { - const std::vector>& new_files = - increment.GetNewFilesIncrement().NewFiles(); - if (!new_files.empty()) { - realtime_context_->AdvanceMaterializedMaxSequenceNumber( - partition_bucket_, DataFileMeta::GetMaxSequenceNumber(new_files)); - } - increment.SetRealtimeOffsetRange(segment.value()->GetOffsetRange()); + increment.SetRealtimeOffsetRange(sealed_range.value()); } return increment; } -Status RealtimePrimaryKeyWriter::FlushSegment( - const std::shared_ptr& segment) { +Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr& segment, + int64_t expected_raw_row_count) { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, realtime_store_->CreateCommitReaders(segment)); ScopeGuard readers_guard([&readers]() { @@ -124,72 +282,26 @@ Status RealtimePrimaryKeyWriter::FlushSegment( } } }); - for (const std::unique_ptr& reader : readers) { + int64_t raw_row_count = 0; + std::vector> sorted_readers; + sorted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { if (!reader) { return Status::Invalid("PK real-time store returned a null commit reader"); } + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr prepared_reader, + AdaptPreparedBatchReader(std::move(reader), prepared_schema_, std::nullopt, key_schema_, + write_schema_, memory_pool_, &raw_row_count)); + auto merge_function = std::make_unique(/*ignore_delete=*/false); + sorted_readers.push_back(std::make_unique( + std::move(prepared_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); } - ConcatBatchReader reader(std::move(readers), memory_pool_); - ScopeGuard reader_guard([&reader]() { reader.Close(); }); - const OffsetRange offset_range = segment->GetOffsetRange(); - int64_t emitted_rows = 0; - while (true) { - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader.NextBatch()); - if (BatchReader::IsEofBatch(batch)) { - break; - } - auto& [c_array, c_schema] = batch; - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, - arrow::ImportArray(c_array.get(), c_schema.get())); - if (!imported || imported->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("PK real-time store commit reader returned a non-StructArray"); - } - std::shared_ptr struct_array = - checked_pointer_cast(imported); - std::shared_ptr value_kind = - struct_array->GetFieldByName(SpecialFields::ValueKind().Name()); - if (!value_kind || value_kind->type_id() != arrow::Type::INT8) { - return Status::Invalid( - "PK real-time store commit reader must return an INT8 _VALUE_KIND field"); - } - std::shared_ptr encoded_row_kinds = - checked_pointer_cast(value_kind); - std::vector row_kinds; - row_kinds.reserve(static_cast(encoded_row_kinds->length())); - for (int64_t i = 0; i < encoded_row_kinds->length(); ++i) { - if (encoded_row_kinds->IsNull(i)) { - return Status::Invalid("PK real-time store commit reader returned a null row kind"); - } - PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, - RowKind::FromByteValue(encoded_row_kinds->Value(i))); - row_kinds.push_back(static_cast(row_kind->ToByteValue())); - } - PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( - struct_array, SpecialFields::ValueKind().Name())); - if (!struct_array->type()->Equals(arrow::struct_(write_schema_->fields()))) { - return Status::Invalid( - "PK real-time store commit reader schema does not match table write schema"); - } - const int64_t row_count = struct_array->length(); - if (row_count > offset_range.Count() - emitted_rows) { - return Status::Invalid( - "PK real-time store commit readers returned more rows than the sealed offset " - "range"); - } - emitted_rows += row_count; - if (row_count == 0) { - continue; - } - auto output = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, output.get())); - RecordBatchBuilder builder(output.get()); - builder.SetRowKinds(row_kinds); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr record_batch, builder.Finish()); - PAIMON_RETURN_NOT_OK(merge_tree_writer_->Write(std::move(record_batch))); - } - if (emitted_rows != offset_range.Count()) { - return Status::Invalid( - "PK real-time store commit readers returned fewer rows than the sealed offset range"); + readers_guard.Release(); + PAIMON_RETURN_NOT_OK(merge_tree_writer_->WriteSortedReaders(std::move(sorted_readers))); + if (raw_row_count != expected_raw_row_count) { + return Status::Invalid("PK real-time store commit readers did not cover the sealed range"); } return Status::OK(); } @@ -197,27 +309,21 @@ Status RealtimePrimaryKeyWriter::FlushSegment( Status RealtimePrimaryKeyWriter::Compact(bool) { return Status::Invalid("PK real-time write does not support explicit compaction"); } - uint64_t RealtimePrimaryKeyWriter::GetMemoryUsage() const { return realtime_store_->GetMemoryUsage(); } - Status RealtimePrimaryKeyWriter::FlushMemory() { return Status::OK(); } - Result RealtimePrimaryKeyWriter::CompactNotCompleted() { return merge_tree_writer_->CompactNotCompleted(); } - Status RealtimePrimaryKeyWriter::Sync() { return merge_tree_writer_->Sync(); } - Status RealtimePrimaryKeyWriter::Close() { return merge_tree_writer_->Close(); } - std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { return merge_tree_writer_->GetMetrics(); } diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index c1e893c85..6abb1ccd0 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -20,16 +20,16 @@ #pragma once #include -#include #include #include #include +#include #include "paimon/core/utils/batch_writer.h" -#include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" namespace arrow { +class MemoryPool; class Schema; } // namespace arrow @@ -37,18 +37,19 @@ namespace paimon { class MemoryPool; class MergeTreeWriter; -class RealtimeContextImpl; +class FieldsComparator; struct RealtimeStoreState; -/// Primary-key real-time writer backed by an in-memory mutation indexer. +/// Coordinates framework-prepared primary-key real-time writes. class RealtimePrimaryKeyWriter final : public BatchWriter { public: static Result> Create( - const std::map& partition, int32_t bucket, const std::shared_ptr& write_schema, - const std::shared_ptr& realtime_context, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const RealtimeStoreState& store_state, int64_t restore_max_sequence_number, const std::shared_ptr& merge_tree_writer, - const std::shared_ptr& memory_pool, const RealtimeStoreState& store_state); + const std::shared_ptr& memory_pool); Status Write(std::unique_ptr&& batch) override; Result PrepareCommit(bool wait_compaction) override; @@ -63,20 +64,28 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { private: RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, const std::shared_ptr& merge_tree_writer, - const std::shared_ptr& realtime_context, - const RealtimePartitionBucket& partition_bucket, const std::shared_ptr& write_schema, - int64_t next_offset, const std::shared_ptr& memory_pool); + const std::shared_ptr& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + int64_t next_offset, int64_t last_sequence_number, + const std::shared_ptr& memory_pool); - Status FlushSegment(const std::shared_ptr& segment); + Status FlushSegment(const std::shared_ptr& segment, + int64_t expected_raw_row_count); std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; std::shared_ptr realtime_store_; std::shared_ptr merge_tree_writer_; - std::shared_ptr realtime_context_; - RealtimePartitionBucket partition_bucket_; std::shared_ptr write_schema_; + std::shared_ptr prepared_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; int64_t next_offset_; + int64_t last_sequence_number_; std::mutex realtime_store_mutex_; std::mutex prepare_mutex_; }; diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 76160ac93..f510c987e 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -24,20 +24,21 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" -#include "paimon/common/data/columnar/columnar_row_ref.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" -#include "paimon/common/types/row_kind.h" -#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/io/prepared_key_value_reader.h" #include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" -#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_fields.h" #include "paimon/core/realtime/realtime_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" @@ -54,177 +55,60 @@ struct ColumnarBatchContext; namespace { -class QueryBatchKeyValueReader final : public KeyValueRecordReader { - public: - QueryBatchKeyValueReader(std::unique_ptr&& reader, - const std::shared_ptr& key_schema, - const std::shared_ptr& value_schema, - const std::shared_ptr& pool) - : reader_(std::move(reader)), - key_schema_(key_schema), - value_schema_(value_schema), - pool_(pool) {} - - ~QueryBatchKeyValueReader() override { - Close(); - } - - Result> NextBatch() override; - std::shared_ptr GetReaderMetrics() const override; - void Close() override; - - private: - class Iterator; - - std::unique_ptr reader_; - std::shared_ptr key_schema_; - std::shared_ptr value_schema_; - std::shared_ptr pool_; - std::shared_ptr values_; - std::shared_ptr sequences_; - std::shared_ptr row_kinds_; - std::shared_ptr key_context_; - std::shared_ptr value_context_; - bool closed_ = false; -}; - -class QueryBatchKeyValueReader::Iterator final : public KeyValueRecordReader::Iterator { - public: - explicit Iterator(QueryBatchKeyValueReader* reader) : reader_(reader) {} - - Result HasNext() const override { - return cursor_ < reader_->values_->length(); - } - - Result Next() override { - if (reader_->sequences_->IsNull(cursor_) || reader_->row_kinds_->IsNull(cursor_)) { - return Status::Invalid("PK merge metadata must not be null"); - } - PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, - RowKind::FromByteValue(reader_->row_kinds_->Value(cursor_))); - const int64_t sequence = reader_->sequences_->Value(cursor_); - std::shared_ptr key = - std::make_shared(reader_->key_context_, cursor_); - auto value = std::make_unique(reader_->value_context_, cursor_++); - return KeyValue(row_kind, sequence, KeyValue::UNKNOWN_LEVEL, std::move(key), - std::move(value)); - } - - private: - QueryBatchKeyValueReader* reader_; - int64_t cursor_ = 0; -}; - -Result> QueryBatchKeyValueReader::NextBatch() { - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); - if (BatchReader::IsEofBatch(batch)) { - return std::unique_ptr(); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, - arrow::ImportArray(batch.first.get(), batch.second.get())); - std::shared_ptr input = - std::dynamic_pointer_cast(imported); - if (!input) { - return Status::Invalid("PK merge input is not a StructArray"); - } - sequences_ = std::dynamic_pointer_cast( - input->GetFieldByName(SpecialFields::SequenceNumber().Name())); - row_kinds_ = std::dynamic_pointer_cast( - input->GetFieldByName(SpecialFields::ValueKind().Name())); - if (!sequences_ || !row_kinds_) { - return Status::Invalid("PK merge input is missing sequence or value-kind metadata"); - } - PAIMON_ASSIGN_OR_RAISE(input, ArrowUtils::RemoveFieldFromStructArray( - input, SpecialFields::SequenceNumber().Name())); - PAIMON_ASSIGN_OR_RAISE( - values_, ArrowUtils::RemoveFieldFromStructArray(input, SpecialFields::ValueKind().Name())); - if (!ArrowUtils::EqualsIgnoreNullable(values_->type(), - arrow::struct_(value_schema_->fields()))) { - return Status::Invalid("PK merge input value schema does not match the table read schema"); - } - arrow::ArrayVector key_fields; - key_fields.reserve(key_schema_->num_fields()); - for (const std::shared_ptr& field : key_schema_->fields()) { - std::shared_ptr key = values_->GetFieldByName(field->name()); - if (!key) { - return Status::Invalid("PK merge input is missing key field ", field->name()); - } - key_fields.push_back(std::move(key)); - } - key_context_ = std::make_shared(key_fields, pool_); - value_context_ = std::make_shared(values_->fields(), pool_); - return std::make_unique(this); -} - -std::shared_ptr QueryBatchKeyValueReader::GetReaderMetrics() const { - return reader_->GetReaderMetrics(); -} - -void QueryBatchKeyValueReader::Close() { - if (closed_) { - return; - } - closed_ = true; - values_.reset(); - sequences_.reset(); - row_kinds_.reset(); - key_context_.reset(); - value_context_.reset(); - if (reader_) { - reader_->Close(); - } -} - Result> CreateMemoryReaders( const std::shared_ptr& split, const RealtimePartitionBucketView& memory, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, const std::shared_ptr& context, const std::shared_ptr& memory_pool) { - arrow::FieldVector requested_fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; - requested_fields.insert(requested_fields.end(), value_schema->fields().begin(), - value_schema->fields().end()); + std::shared_ptr full_value_schema = + DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields()); + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), full_value_schema->fields().begin(), + full_value_schema->fields().end()); + std::shared_ptr prepared_schema = arrow::schema(std::move(prepared_fields)); auto c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportSchema(*arrow::schema(requested_fields), c_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema, c_schema.get())); ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); - RealtimeQueryContext query_context{c_schema.get(), /*predicate=*/nullptr, - /*enable_predicate_pushdown=*/false}; - PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, - memory.store->CreateQueryReaders( - memory.read_view, split->CommittedEndOffset(), query_context)); - ScopeGuard reader_guard([&batch_readers]() { + RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> batch_readers, + memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); + ScopeGuard batch_readers_guard([&batch_readers]() { for (const std::unique_ptr& reader : batch_readers) { if (reader) { reader->Close(); } } }); - if (batch_readers.empty()) { - return Status::Invalid("PK real-time store returned no query readers for active memory"); - } std::vector result; result.reserve(batch_readers.size()); for (std::unique_ptr& reader : batch_readers) { if (!reader) { return Status::Invalid("PK real-time store returned a null query reader"); } - std::shared_ptr min_key; - std::shared_ptr max_key; - if (auto* provider = dynamic_cast(reader.get())) { - min_key = provider->GetMinKey(); - max_key = provider->GetMaxKey(); - } - result.push_back( - AdditionalKeyValueReader{std::make_unique( - std::move(reader), key_schema, value_schema, memory_pool), - std::move(min_key), std::move(max_key)}); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_reader, + AdaptPreparedBatchReader(std::move(reader), prepared_schema, + OffsetRange(split->CommittedEndOffset(), + split->MemoryEndOffset()), + key_schema, value_schema, memory_pool)); + auto merge = std::make_unique(false); + result.push_back(AdditionalKeyValueReader{ + std::make_unique( + std::move(prepared_reader), key_comparator, + std::make_shared(std::move(merge))), + nullptr, nullptr}); } + batch_readers_guard.Release(); return result; } -} // namespace +} KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, @@ -268,7 +152,7 @@ Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); if (realtime_split) { - return CreateRealtimeReader(realtime_split, /*release_ticket=*/true); + return CreateRealtimeReader(realtime_split, true); } std::shared_ptr dispatch_split = split; @@ -332,7 +216,7 @@ Result> KeyValueTableRead::CreateReader( std::dynamic_pointer_cast(split); if (realtime_split) { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, - CreateRealtimeReader(realtime_split, /*release_ticket=*/false)); + CreateRealtimeReader(realtime_split, false)); readers.push_back(std::move(reader)); realtime_splits.push_back(std::move(realtime_split)); } else { @@ -386,7 +270,8 @@ Result> KeyValueTableRead::CreateRealtimeReader( PAIMON_ASSIGN_OR_RAISE( std::vector memory_readers, CreateMemoryReaders(realtime_split, memory, merge_read->GetKeySchema(), - merge_read->GetValueSchema(), context_, GetMemoryPool())); + merge_read->GetValueSchema(), merge_read->GetKeyComparator(), + context_, GetMemoryPool())); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, merge_read->CreateRealtimeReader(realtime_split->DiskSplits(), std::move(memory_readers))); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 9f302eb37..be7595381 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -443,9 +443,52 @@ class CloseTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { std::shared_ptr state_; }; -class InvalidReaderRealtimeStore final : public RealtimeStore { +class SplitBatchReader final : public BatchReader { public: - explicit InvalidReaderRealtimeStore(const std::shared_ptr& delegate) + explicit SplitBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + while (!current_batch_ || next_row_ == current_batch_->length()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("split batch reader received a non-struct batch"); + } + current_batch_ = std::dynamic_pointer_cast(array); + next_row_ = 0; + } + std::shared_ptr slice = current_batch_->Slice(next_row_, /*length=*/1); + ++next_row_; + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*slice, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + current_batch_.reset(); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr current_batch_; + int64_t next_row_ = 0; +}; + +class SplitCommitReaderRealtimeStore final : public RealtimeStore { + public: + explicit SplitCommitReaderRealtimeStore(const std::shared_ptr& delegate) : delegate_(delegate) {} Status Write(RealtimeWriteBatch&& batch) override { @@ -457,9 +500,12 @@ class InvalidReaderRealtimeStore final : public RealtimeStore { } Result>> CreateCommitReaders( - const std::shared_ptr&) override { - std::vector> readers; - readers.push_back(nullptr); + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader)); + } return readers; } @@ -468,8 +514,9 @@ class InvalidReaderRealtimeStore final : public RealtimeStore { } Result>> CreateQueryReaders( - const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { - return std::vector>(); + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); } Status AdvanceCommittedOffset(int64_t committed_offset) override { @@ -484,13 +531,13 @@ class InvalidReaderRealtimeStore final : public RealtimeStore { std::shared_ptr delegate_; }; -class InvalidReaderRealtimeStoreFactory final : public RealtimeStoreFactory { +class SplitCommitReaderRealtimeStoreFactory final : public RealtimeStoreFactory { public: Result> Create(RealtimeStoreCreateRequest&& request) override { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, delegate_.Create(std::move(request))); return std::shared_ptr( - std::make_shared(delegate)); + std::make_shared(delegate)); } private: @@ -1311,10 +1358,11 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter(realtime_context)); - std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}}; + std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}, {1, "new-in-run", "p0"}}; ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, MakeBatch(first_rows, /*partitioned=*/false, /*bucket=*/0, - {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT})); + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); ASSERT_OK(writer->Write(std::move(first_batch))); ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, MakeBatch({Row{1, "new", "p0"}}, /*partitioned=*/false, /*bucket=*/0, @@ -1756,13 +1804,13 @@ TEST_F(RealtimeWriteInteTest, TestPkRecovery) { MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); ASSERT_OK(first_writer->Write(std::move(batch))); ASSERT_OK_AND_ASSIGN(std::vector memory_sequences, ReadPkSequences(first_context)); - ASSERT_EQ((std::vector{2, 3, 4}), memory_sequences); + ASSERT_EQ((std::vector{1, 2, 3, 4}), memory_sequences); ASSERT_OK_AND_ASSIGN(std::vector progress, first_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); ASSERT_EQ(1, progress.size()); ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); ASSERT_EQ(1, NewFiles(progress).size()); - ASSERT_EQ(memory_sequences.front(), NewFiles(progress)[0]->min_sequence_number); + ASSERT_EQ(2, NewFiles(progress)[0]->min_sequence_number); ASSERT_EQ(memory_sequences.back(), NewFiles(progress)[0]->max_sequence_number); ASSERT_OK(Commit(progress, /*commit_identifier=*/1)); ASSERT_OK(first_writer->Close()); @@ -1801,9 +1849,16 @@ TEST_F(RealtimeWriteInteTest, TestPkCompaction) { constexpr int64_t kCommitRoundsBeforeCompaction = 4; std::set committed_file_names; for (int64_t round = 0; round < kCommitRoundsBeforeCompaction; ++round) { - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch({Row{round, "value-" + std::to_string(round), "p0"}}, - /*partitioned=*/false)); + const bool delete_latest_live_row = round == kCommitRoundsBeforeCompaction - 1; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeBatch( + {Row{delete_latest_live_row ? round - 1 : round, + delete_latest_live_row ? "deleted" : "value-" + std::to_string(round), "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + delete_latest_live_row + ? std::vector{RecordBatch::RowKind::DELETE} + : std::vector{})); ASSERT_OK(writer->Write(std::move(batch))); ASSERT_OK_AND_ASSIGN(std::vector progress, writer->PrepareCommitWithProgress(round)); @@ -1819,11 +1874,6 @@ TEST_F(RealtimeWriteInteTest, TestPkCompaction) { ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); ASSERT_EQ(0, memory_usage); } - ASSERT_OK_AND_ASSIGN(std::unique_ptr next_batch, - MakeBatch({Row{4, "value-4", "p0"}}, - /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(next_batch))); - WriteContextBuilder compact_builder(table_path_, commit_user_); compact_builder.SetOptions(options_).WithStreamingMode(true); ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_context, compact_builder.Finish()); @@ -1848,6 +1898,14 @@ TEST_F(RealtimeWriteInteTest, TestPkCompaction) { } ASSERT_EQ(committed_file_names, compacted_file_names); ASSERT_FALSE(compact_message->GetCompactIncrement().CompactAfter().empty()); + constexpr int64_t kHistoricalMaxSequenceNumber = kCommitRoundsBeforeCompaction - 1; + int64_t compacted_live_max_sequence_number = -1; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactAfter()) { + compacted_live_max_sequence_number = + std::max(compacted_live_max_sequence_number, file->max_sequence_number); + } + ASSERT_LT(compacted_live_max_sequence_number, kHistoricalMaxSequenceNumber); ASSERT_OK(CommitMessages(compact_messages, /*commit_identifier=*/4)); ASSERT_OK(compact_writer->Close()); @@ -1859,45 +1917,38 @@ TEST_F(RealtimeWriteInteTest, TestPkCompaction) { ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot->GetCommitKind()); ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); ASSERT_EQ(4, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); - ASSERT_OK(writer->RefreshCommittedSnapshot(compact_snapshot->Id())); - ASSERT_OK_AND_ASSIGN(std::vector compacted_rows, ReadRows(realtime_context)); - ASSERT_EQ((std::vector{{0, "value-0", "p0"}, - {1, "value-1", "p0"}, - {2, "value-2", "p0"}, - {3, "value-3", "p0"}, - {4, "value-4", "p0"}}), - compacted_rows); - - constexpr int64_t kCommitRoundsAfterCompaction = 2; - for (int64_t round = 0; round < kCommitRoundsAfterCompaction; ++round) { - if (round > 0) { - ASSERT_OK_AND_ASSIGN( - std::unique_ptr batch, - MakeBatch({Row{4 + round, "value-" + std::to_string(4 + round), "p0"}}, - /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(batch))); - } - const int64_t commit_identifier = 5 + round; - ASSERT_OK_AND_ASSIGN(std::vector progress, - writer->PrepareCommitWithProgress(commit_identifier)); - ASSERT_EQ(1, progress.size()); - ASSERT_EQ(OffsetRange(4 + round, 5 + round), progress[0].offset_range); - ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(progress, commit_identifier)); - ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); - ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); - ASSERT_EQ(0, memory_usage); - } + ASSERT_OK_AND_ASSIGN(std::vector compacted_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}}), compacted_rows); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr fresh_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_writer, + CreateRealtimeWriter(fresh_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_batch, + MakeBatch({Row{4, "value-4", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(fresh_writer->Write(std::move(fresh_batch))); + ASSERT_OK_AND_ASSIGN(std::vector fresh_sequences, ReadPkSequences(fresh_context)); + ASSERT_EQ((std::vector{compacted_live_max_sequence_number + 1}), fresh_sequences); + ASSERT_LT(fresh_sequences.front(), kHistoricalMaxSequenceNumber); + ASSERT_OK_AND_ASSIGN(std::vector fresh_progress, + fresh_writer->PrepareCommitWithProgress(/*commit_identifier=*/5)); + ASSERT_EQ(1, fresh_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), fresh_progress[0].offset_range); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->min_sequence_number); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->max_sequence_number); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(fresh_progress, /*commit_identifier=*/5)); + ASSERT_OK(fresh_writer->Close()); + ASSERT_OK_AND_ASSIGN(offsets, ReadCommittedOffsets()); - ASSERT_EQ(6, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); - ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); - ASSERT_EQ((std::vector{{0, "value-0", "p0"}, - {1, "value-1", "p0"}, - {2, "value-2", "p0"}, - {3, "value-3", "p0"}, - {4, "value-4", "p0"}, - {5, "value-5", "p0"}}), + ASSERT_EQ(5, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}, {4, "value-4", "p0"}}), final_rows); - ASSERT_OK(writer->Close()); } TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { @@ -2054,19 +2105,29 @@ TEST_F(RealtimeWriteInteTest, TestPkWriteDuringPrepare) { TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { CreatePkTable(); - auto factory = std::make_shared(); + auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create(factory)); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter(realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{4, "four", "p0"}, Row{3, "three", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch({Row{2, "two", "p0"}, Row{1, "one", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); - ASSERT_NOK_WITH_MSG(ReadRows(realtime_context), - "PK real-time store returned no query readers for active memory"); - ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), - "PK real-time store returned a null commit reader"); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector rows, ReadRows()); + ASSERT_EQ((std::vector{ + {1, "one", "p0"}, {2, "two", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}), + rows); ASSERT_OK(writer->Close()); } @@ -2757,52 +2818,6 @@ TEST_F(RealtimeWriteInteTest, TestCloseWriterAllowsContextReuseByLaterWriter) { ASSERT_OK(second_writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { - CreatePkTable(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, - CreateRealtimeWriter(realtime_context)); - std::vector first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, - MakeBatch(first_rows, /*partitioned=*/false)); - ASSERT_OK(first_writer->Write(std::move(first_batch))); - ASSERT_OK_AND_ASSIGN(std::vector commits, - first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); - ASSERT_EQ(1, commits.size()); - ASSERT_EQ(OffsetRange(0, 3), commits[0].offset_range); - ASSERT_EQ(1, NewFiles(commits).size()); - ASSERT_EQ(0, NewFiles(commits)[0]->min_sequence_number); - ASSERT_EQ(2, NewFiles(commits)[0]->max_sequence_number); - ASSERT_OK(first_writer->Close()); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, - CreateRealtimeWriter(realtime_context)); - std::vector second_rows = { - Row{0, "updated-0", "p0"}, - Row{3, "value-3", "p0"}, - }; - ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, - MakeBatch(second_rows, /*partitioned=*/false)); - ASSERT_OK(second_writer->Write(std::move(second_batch))); - ASSERT_OK_AND_ASSIGN(std::vector second_commits, - second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_EQ(1, second_commits.size()); - ASSERT_EQ(OffsetRange(3, 5), second_commits[0].offset_range); - ASSERT_EQ(1, NewFiles(second_commits).size()); - ASSERT_EQ(3, NewFiles(second_commits)[0]->min_sequence_number); - ASSERT_EQ(4, NewFiles(second_commits)[0]->max_sequence_number); - - commits.push_back(std::move(second_commits[0])); - ASSERT_OK(Commit(commits, /*commit_identifier=*/1)); - std::vector expected_rows = first_rows; - expected_rows[0] = second_rows[0]; - expected_rows.push_back(second_rows[1]); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); - ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(second_writer->Close()); -} - TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, From ea90f89236bb8869aeb91f0ad804ccbb0e5b8b84 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:35:37 +0800 Subject: [PATCH 16/24] refactor(realtime): simplify primary-key write preparation --- include/paimon/realtime/realtime_context.h | 4 - src/paimon/CMakeLists.txt | 2 +- .../merged_key_value_record_reader_test.cpp | 2 +- .../core/mergetree/merge_tree_writer_test.cpp | 15 - src/paimon/core/mergetree/write_buffer.cpp | 4 - .../key_value_file_store_write_test.cpp | 98 ------- .../prepared_key_value_reader.cpp | 2 +- .../prepared_key_value_reader.h | 0 .../realtime/realtime_primary_key_writer.cpp | 2 +- .../table/source/key_value_table_read.cpp | 2 +- .../core/utils/primary_key_table_utils.h | 1 - test/inte/realtime_write_inte_test.cpp | 273 ------------------ 12 files changed, 5 insertions(+), 400 deletions(-) rename src/paimon/core/{io => realtime}/prepared_key_value_reader.cpp (99%) rename src/paimon/core/{io => realtime}/prepared_key_value_reader.h (100%) diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index 8f2967b32..200e4ba4c 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -78,10 +78,6 @@ using RealtimeOffsetMap = std::map; /// partition drop, and rollback operations do not automatically clear process-local real-time /// state. Applications must coordinate these operations with active real-time writers and recreate /// the `RealtimeContext` and writers before continuing. -/// -/// A primary-key writer and its context form one lifecycle. After a primary-key write or prepare -/// returns an error, discard both, create fresh instances from the latest committed snapshot, and -/// replay the caller-owned WAL. Reusing the failed context or retrying prepare is unsupported. class PAIMON_EXPORT RealtimeContext { public: /// Creates a context backed by Paimon's default in-memory Arrow `RealtimeStore`. diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index fc1fb00dd..49871672d 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -282,7 +282,6 @@ set(PAIMON_CORE_SRCS core/io/vector_file_batch_reader.cpp core/io/file_index_evaluator.cpp core/io/key_value_data_file_record_reader.cpp - core/io/prepared_key_value_reader.cpp core/io/key_value_data_file_writer_factory.cpp core/io/key_value_data_file_writer.cpp core/io/key_value_in_memory_record_reader.cpp @@ -379,6 +378,7 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/prepared_key_value_reader.cpp core/realtime/primary_key_realtime_store.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 39714fa29..21b0a16b1 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -34,9 +34,9 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" -#include "paimon/core/io/prepared_key_value_reader.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include "paimon/core/realtime/realtime_fields.h" #include "paimon/memory/memory_pool.h" #include "paimon/realtime/offset_range.h" diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index aa2d0c959..675ce3198 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include @@ -613,20 +612,6 @@ TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { ASSERT_OK(merge_writer->Close()); } -TEST_P(MergeTreeWriterTest, TestRejectsExhaustedSequence) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, - CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); - - auto dir = UniqueTestDirectory::Create(); - ASSERT_TRUE(dir); - auto path_factory = std::make_shared(); - ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); - - ASSERT_NOK_WITH_MSG(CreateMergeWriter(std::numeric_limits::max(), dir->Str(), - path_factory, 0, options), - "sequence number has reached INT64_MAX"); -} - TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({ diff --git a/src/paimon/core/mergetree/write_buffer.cpp b/src/paimon/core/mergetree/write_buffer.cpp index 3d3fdc196..549975a33 100644 --- a/src/paimon/core/mergetree/write_buffer.cpp +++ b/src/paimon/core/mergetree/write_buffer.cpp @@ -18,7 +18,6 @@ #include "paimon/core/mergetree/write_buffer.h" -#include #include #include @@ -40,9 +39,6 @@ Result> WriteBuffer::Create( const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, const std::shared_ptr& io_manager, bool enable_multi_thread_spill, const std::shared_ptr& pool) { - if (last_sequence_number == std::numeric_limits::max()) { - return Status::Invalid("sequence number has reached INT64_MAX"); - } auto value_type = arrow::struct_(value_schema->fields()); auto in_memory_buffer = std::make_unique( last_sequence_number, value_type, trimmed_primary_keys, user_defined_sequence_fields, diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index cbd2189fc..733c19d6e 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -59,7 +59,6 @@ #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" -#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" @@ -111,69 +110,6 @@ class TestingMemoryPool final : public MemoryPool { std::unique_ptr delegate_ = GetMemoryPool(); }; -class FailOnceRealtimeStore final : public RealtimeStore { - public: - FailOnceRealtimeStore(const std::shared_ptr& delegate, - const std::shared_ptr& fail_next_write) - : delegate_(delegate), fail_next_write_(fail_next_write) {} - - Status Write(RealtimeWriteBatch&& batch) override { - if (*fail_next_write_) { - *fail_next_write_ = false; - return Status::Invalid("injected real-time store write failure"); - } - return delegate_->Write(std::move(batch)); - } - - Result>> SealForCommit() override { - return delegate_->SealForCommit(); - } - - Result>> CreateCommitReaders( - const std::shared_ptr& segment) override { - return delegate_->CreateCommitReaders(segment); - } - - Result> AcquireReadView() override { - return delegate_->AcquireReadView(); - } - - Result>> CreateQueryReaders( - const std::shared_ptr& view, int64_t offset_begin, - const RealtimeQueryContext& context) override { - return delegate_->CreateQueryReaders(view, offset_begin, context); - } - - Status AdvanceCommittedOffset(int64_t committed_offset) override { - return delegate_->AdvanceCommittedOffset(committed_offset); - } - - uint64_t GetMemoryUsage() const override { - return delegate_->GetMemoryUsage(); - } - - private: - std::shared_ptr delegate_; - std::shared_ptr fail_next_write_; -}; - -class FailOnceRealtimeStoreFactory final : public RealtimeStoreFactory { - public: - explicit FailOnceRealtimeStoreFactory(const std::shared_ptr& fail_next_write) - : fail_next_write_(fail_next_write) {} - - Result> Create(RealtimeStoreCreateRequest&& request) override { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, - delegate_.Create(std::move(request))); - return std::shared_ptr( - std::make_shared(delegate, fail_next_write_)); - } - - private: - ArrowRealtimeStoreFactory delegate_; - std::shared_ptr fail_next_write_; -}; - } class KeyValueFileStoreWriteTest : public ::testing::Test { @@ -551,40 +487,6 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeOffsetCollision) { ASSERT_OK(writer->Close()); } -TEST_F(KeyValueFileStoreWriteTest, TestWriteFailureKeepsCursors) { - const std::map options = { - {Options::BUCKET, "1"}, - {Options::WRITE_BUFFER_SIZE, "1"}, - }; - const std::shared_ptr schema = arrow::schema({ - arrow::field("id", arrow::int64(), false), - arrow::field("value", arrow::utf8()), - }); - std::unique_ptr dir = UniqueTestDirectory::Create(); - ASSERT_TRUE(dir); - CreateTable(dir->Str(), schema, options); - const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); - - auto fail_next_write = std::make_shared(true); - auto factory = std::make_shared(fail_next_write); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create(factory)); - WriteContextBuilder builder(table_path, "test"); - builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); - ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - FileStoreWrite::Create(std::move(write_context))); - - ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[9, "rejected"]])")), - "injected real-time store write failure"); - ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "kept"]])"))); - using PreparedRow = std::tuple; - ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, - ReadPreparedRows(realtime_context)); - ASSERT_EQ((std::vector{{0, 1, "kept", 0, 0}}), prepared_rows); - ASSERT_OK(writer->Close()); -} - TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { const std::map options = {{Options::BUCKET, "1"}}; const std::shared_ptr schema = arrow::schema({ diff --git a/src/paimon/core/io/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp similarity index 99% rename from src/paimon/core/io/prepared_key_value_reader.cpp rename to src/paimon/core/realtime/prepared_key_value_reader.cpp index 0f4f22097..b99f67dd9 100644 --- a/src/paimon/core/io/prepared_key_value_reader.cpp +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -16,7 +16,7 @@ * limitations under the License. */ -#include "paimon/core/io/prepared_key_value_reader.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include #include diff --git a/src/paimon/core/io/prepared_key_value_reader.h b/src/paimon/core/realtime/prepared_key_value_reader.h similarity index 100% rename from src/paimon/core/io/prepared_key_value_reader.h rename to src/paimon/core/realtime/prepared_key_value_reader.h diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index c85ff6322..2dc5a71b4 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -34,10 +34,10 @@ #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/merged_key_value_record_reader.h" -#include "paimon/core/io/prepared_key_value_reader.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/realtime/realtime_fields.h" #include "paimon/core/utils/commit_increment.h" diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index f510c987e..31779e049 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -31,12 +31,12 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/merged_key_value_record_reader.h" -#include "paimon/core/io/prepared_key_value_reader.h" #include "paimon/core/key_value.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/realtime/realtime_fields.h" #include "paimon/core/realtime/realtime_reader.h" diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index c40e92cda..82a108ab7 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -24,7 +24,6 @@ #include "arrow/type.h" #include "paimon/result.h" -#include "paimon/status.h" namespace arrow { class Schema; diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index be7595381..c61b04b0d 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -23,7 +23,6 @@ #include #include #include -#include #include #include #include @@ -79,126 +78,6 @@ namespace paimon::test { namespace { -class BlockingState { - public: - void Block() { - std::unique_lock lock(mutex_); - entered_ = true; - entered_cv_.notify_all(); - release_cv_.wait(lock, [this]() { return released_; }); - } - - bool WaitUntilBlocked() { - std::unique_lock lock(mutex_); - return entered_cv_.wait_for(lock, std::chrono::seconds(30), [this]() { return entered_; }); - } - - void Release() { - std::lock_guard lock(mutex_); - released_ = true; - release_cv_.notify_all(); - } - - private: - std::mutex mutex_; - std::condition_variable entered_cv_; - std::condition_variable release_cv_; - bool entered_ = false; - bool released_ = false; -}; - -class BlockingBatchReader final : public BatchReader { - public: - BlockingBatchReader(std::unique_ptr&& reader, - const std::shared_ptr& state) - : reader_(std::move(reader)), state_(state) {} - - Result NextBatch() override { - if (!blocked_) { - blocked_ = true; - state_->Block(); - } - return reader_->NextBatch(); - } - - std::shared_ptr GetReaderMetrics() const override { - return reader_->GetReaderMetrics(); - } - - void Close() override { - reader_->Close(); - } - - private: - std::unique_ptr reader_; - std::shared_ptr state_; - bool blocked_ = false; -}; - -class BlockingRealtimeStore final : public RealtimeStore { - public: - BlockingRealtimeStore(const std::shared_ptr& delegate, - const std::shared_ptr& state) - : delegate_(delegate), state_(state) {} - - Status Write(RealtimeWriteBatch&& batch) override { - return delegate_->Write(std::move(batch)); - } - - Result>> SealForCommit() override { - return delegate_->SealForCommit(); - } - - Result>> CreateCommitReaders( - const std::shared_ptr& segment) override { - PAIMON_ASSIGN_OR_RAISE(std::vector> readers, - delegate_->CreateCommitReaders(segment)); - if (!readers.empty()) { - readers[0] = std::make_unique(std::move(readers[0]), state_); - } - return readers; - } - - Result> AcquireReadView() override { - return delegate_->AcquireReadView(); - } - - Result>> CreateQueryReaders( - const std::shared_ptr& view, int64_t offset_begin, - const RealtimeQueryContext& context) override { - return delegate_->CreateQueryReaders(view, offset_begin, context); - } - - Status AdvanceCommittedOffset(int64_t committed_offset) override { - return delegate_->AdvanceCommittedOffset(committed_offset); - } - - uint64_t GetMemoryUsage() const override { - return delegate_->GetMemoryUsage(); - } - - private: - std::shared_ptr delegate_; - std::shared_ptr state_; -}; - -class BlockingRealtimeStoreFactory final : public RealtimeStoreFactory { - public: - explicit BlockingRealtimeStoreFactory(const std::shared_ptr& state) - : state_(state) {} - - Result> Create(RealtimeStoreCreateRequest&& request) override { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, - delegate_.Create(std::move(request))); - return std::shared_ptr( - std::make_shared(delegate, state_)); - } - - private: - ArrowRealtimeStoreFactory delegate_; - std::shared_ptr state_; -}; - class TrackingRealtimeReadView final : public RealtimeReadView { public: explicit TrackingRealtimeReadView(std::shared_ptr delegate) @@ -1951,158 +1830,6 @@ TEST_F(RealtimeWriteInteTest, TestPkCompaction) { final_rows); } -TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { - CreatePkTable(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - - constexpr int64_t kRowCount = 20; - constexpr int32_t kReaderCount = 2; - std::atomic writer_done{false}; - std::atomic control_done{false}; - std::atomic commit_count{0}; - ConcurrentTestState state; - std::vector read_counts(kReaderCount, 0); - - std::thread write_thread([&]() { - state.WaitForStart(); - for (int64_t id = 0; id < kRowCount && !state.ShouldStop(); ++id) { - Result> batch = - MakeBatch(MakeRows(id, /*count=*/1, /*partition=*/"p0"), - /*partitioned=*/false); - if (state.RecordErrorIfNotOk(batch) || - state.RecordErrorIfNotOk(writer->Write(std::move(batch).value()))) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(2)); - } - writer_done.store(true, std::memory_order_release); - }); - - std::thread control_thread([&]() { - state.WaitForStart(); - int64_t commit_identifier = 0; - do { - Result> progress = - writer->PrepareCommitWithProgress(commit_identifier); - if (state.RecordErrorIfNotOk(progress)) { - break; - } - if (!progress.value().empty()) { - Result snapshot = Commit(progress.value(), commit_identifier++); - if (state.RecordErrorIfNotOk(snapshot) || - state.RecordErrorIfNotOk(writer->RefreshCommittedSnapshot(snapshot.value()))) { - break; - } - commit_count.fetch_add(1, std::memory_order_relaxed); - } - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } while (!writer_done.load(std::memory_order_acquire) && !state.ShouldStop()); - if (!state.ShouldStop()) { - Result> progress = - writer->PrepareCommitWithProgress(commit_identifier); - if (!state.RecordErrorIfNotOk(progress) && !progress.value().empty()) { - Result snapshot = Commit(progress.value(), commit_identifier); - if (!state.RecordErrorIfNotOk(snapshot) && - !state.RecordErrorIfNotOk(writer->RefreshCommittedSnapshot(snapshot.value()))) { - commit_count.fetch_add(1, std::memory_order_relaxed); - } - } - } - control_done.store(true, std::memory_order_release); - }); - - std::vector read_threads; - read_threads.reserve(kReaderCount); - for (int32_t reader_index = 0; reader_index < kReaderCount; ++reader_index) { - read_threads.emplace_back([&, reader_index]() { - state.WaitForStart(); - while (!control_done.load(std::memory_order_acquire) && !state.ShouldStop()) { - Result> rows = ReadRows(realtime_context); - ++read_counts[reader_index]; - if (state.RecordErrorIfNotOk(rows) || - state.RecordErrorIfNotOk(ValidateReadPrefix(rows.value(), kRowCount))) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } - }); - } - - state.StartWhenReady(/*worker_count=*/2 + kReaderCount); - write_thread.join(); - control_thread.join(); - for (std::thread& read_thread : read_threads) { - read_thread.join(); - } - - ASSERT_TRUE(state.Errors().empty()) << (state.Errors().empty() ? "" : state.Errors().front()); - ASSERT_GT(commit_count.load(), 0); - for (int32_t read_count : read_counts) { - ASSERT_GT(read_count, 0); - } - ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); - ASSERT_EQ(kRowCount, static_cast(final_rows.size())); - ASSERT_OK(ValidateReadPrefix(final_rows, kRowCount)); - ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); - ASSERT_EQ(0, memory_usage); - ASSERT_OK(writer->Close()); -} - -TEST_F(RealtimeWriteInteTest, TestPkWriteDuringPrepare) { - CreatePkTable(); - auto state = std::make_shared(); - auto factory = std::make_shared(state); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create(factory)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, - MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(first_batch))); - - Result> prepare_result = - Status::Invalid("prepare did not run"); - std::thread prepare_thread( - [&]() { prepare_result = writer->PrepareCommitWithProgress(/*commit_identifier=*/0); }); - const bool prepare_blocked = state->WaitUntilBlocked(); - if (!prepare_blocked) { - state->Release(); - prepare_thread.join(); - ASSERT_TRUE(prepare_blocked); - } - - std::promise write_promise; - std::future write_future = write_promise.get_future(); - std::thread write_thread([&]() { - Result> batch = - MakeBatch({Row{2, "two", "p0"}}, /*partitioned=*/false); - if (!batch.ok()) { - write_promise.set_value(batch.status()); - return; - } - write_promise.set_value(writer->Write(std::move(batch).value())); - }); - const bool write_completed = - write_future.wait_for(std::chrono::seconds(5)) == std::future_status::ready; - state->Release(); - prepare_thread.join(); - write_thread.join(); - - ASSERT_TRUE(write_completed); - ASSERT_OK(write_future.get()); - ASSERT_OK(prepare_result); - ASSERT_EQ(1, prepare_result.value().size()); - ASSERT_EQ(OffsetRange(0, 1), prepare_result.value()[0].offset_range); - ASSERT_OK_AND_ASSIGN(std::vector second_progress, - writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); - ASSERT_EQ(1, second_progress.size()); - ASSERT_EQ(OffsetRange(1, 2), second_progress[0].offset_range); - ASSERT_OK(writer->Close()); -} - TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { CreatePkTable(); auto factory = std::make_shared(); From 53f6b02ef6551630548d67f554067a390a14239f Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:49:07 +0800 Subject: [PATCH 17/24] test(mergetree): reuse reader failure mock --- .../core/mergetree/merge_tree_writer_test.cpp | 44 ++++++------------- 1 file changed, 14 insertions(+), 30 deletions(-) diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 675ce3198..63e896573 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -96,31 +96,7 @@ class TrackingKeyValueRecordReader : public KeyValueRecordReader { bool* closed_flag_; }; -class ErrorKeyValueRecordReader : public KeyValueRecordReader { - public: - ErrorKeyValueRecordReader(Status status, bool* closed_flag) - : status_(std::move(status)), closed_flag_(closed_flag) {} - - Result> NextBatch() override { - return status_; - } - - std::shared_ptr GetReaderMetrics() const override { - return nullptr; - } - - void Close() override { - if (closed_flag_ != nullptr) { - *closed_flag_ = true; - } - } - - private: - Status status_; - bool* closed_flag_; -}; - -} +} // namespace class MergeTreeWriterTest : public ::testing::TestWithParam { public: @@ -270,7 +246,8 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { } std::unique_ptr CreateSingleReader( - const std::shared_ptr& array, int32_t batch_size = 16) const { + const std::shared_ptr& array, int32_t batch_size = 16, + const Status& next_batch_status = Status::OK()) const { std::vector write_fields = {SpecialFields::SequenceNumber(), SpecialFields::ValueKind()}; write_fields.insert(write_fields.end(), value_fields_.begin(), value_fields_.end()); @@ -280,6 +257,7 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { arrow::schema(arrow::FieldVector({write_schema->field(2)})); auto file_batch_reader = std::make_unique(array, array->type(), batch_size); + file_batch_reader->SetNextBatchStatus(next_batch_status); return std::make_unique( std::move(file_batch_reader), key_schema, value_schema_, 0, pool_); } @@ -601,13 +579,19 @@ TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { Status null_status = merge_writer->WriteSortedReaders(std::move(null_readers)); ASSERT_TRUE(null_status.IsInvalid()); + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + Status expected_status = Status::IOError("sorted reader failure"); bool failing_reader_closed = false; - auto failing_reader = std::make_unique( - Status::IOError("sorted reader failure"), &failing_reader_closed); std::vector> failing_readers; - failing_readers.push_back(std::move(failing_reader)); + failing_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array, /*batch_size=*/16, expected_status), + &failing_reader_closed)); Status failing_status = merge_writer->WriteSortedReaders(std::move(failing_readers)); - ASSERT_TRUE(failing_status.IsIOError()); + ASSERT_EQ(expected_status, failing_status); ASSERT_TRUE(failing_reader_closed); ASSERT_OK(merge_writer->Close()); } From 141099c7e4f57e3cdf40d3c865619f61a8b8e7f8 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:49:38 +0800 Subject: [PATCH 18/24] fix(realtime): preserve PK sequence across writer handoff --- .../operation/key_value_file_store_write.cpp | 6 +-- .../core/realtime/realtime_context_impl.cpp | 11 +++++ .../core/realtime/realtime_context_impl.h | 4 ++ .../core/realtime/realtime_context_test.cpp | 16 +++++++ .../realtime/realtime_primary_key_writer.cpp | 23 +++++++--- .../realtime/realtime_primary_key_writer.h | 10 ++++- test/inte/realtime_write_inte_test.cpp | 45 +++++++++++++++++++ 7 files changed, 106 insertions(+), 9 deletions(-) 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 d2c97abcf..de7217ec3 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -175,9 +175,9 @@ Result> KeyValueFileStoreWrite::CreateWriter( if (!realtime_context_) { return std::shared_ptr(std::move(writer)); } - return RealtimePrimaryKeyWriter::Create(schema_, trimmed_primary_keys, key_comparator_, - realtime_store_state.value(), restore_max_seq_number, - writer, pool_); + return RealtimePrimaryKeyWriter::Create( + partition_map, bucket, schema_, trimmed_primary_keys, key_comparator_, + realtime_context_impl, realtime_store_state.value(), restore_max_seq_number, writer, pool_); } Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index b73cfdb8a..415052a69 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -130,6 +130,17 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( return RealtimeStoreState{std::move(store), initial_offset}; } +int64_t RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + std::lock_guard lock(mutex_); + auto [iter, inserted] = + materialized_max_sequence_numbers_.emplace(partition_bucket, max_sequence_number); + if (!inserted && max_sequence_number > iter->second) { + iter->second = max_sequence_number; + } + return iter->second; +} + Result> RealtimeContextImpl::AcquireReadViews() { std::lock_guard lock(mutex_); std::vector result; diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 4f62cf1ee..aa4d263c6 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -67,6 +67,9 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); + int64_t AdvanceMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, + int64_t max_sequence_number); + Result> AcquireReadViews(); Result PinReadView(const RealtimePartitionBucketView& view, int64_t ttl_millis); @@ -96,6 +99,7 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::mutex mutex_; std::mutex progress_mutex_; std::map> stores_; + std::map materialized_max_sequence_numbers_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; // Progress already reflected in stores owned by this context. diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 07bbf555b..15066ca1d 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -158,6 +158,22 @@ TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { ASSERT_EQ(1, factory->stores[2]->acquire_count); } +TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + const RealtimePartitionBucket partition_bucket(partition, /*bucket=*/0); + + ASSERT_EQ(4, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/4)); + ASSERT_EQ(8, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/8)); + ASSERT_EQ(8, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/6)); + ASSERT_EQ(10, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/10)); +} + TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 2dc5a71b4..b53831f0a 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -138,13 +138,16 @@ Result> PrepareBatch( } // namespace Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, const std::shared_ptr& write_schema, const std::vector& trimmed_primary_keys, - const std::shared_ptr& key_comparator, const RealtimeStoreState& store_state, - int64_t restored_max_sequence_number, const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& key_comparator, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, const std::shared_ptr& memory_pool) { if (!store_state.store || !merge_tree_writer || !write_schema || !key_comparator || - !memory_pool) { + !realtime_context || !memory_pool) { return Status::Invalid("PK real-time writer received a null dependency"); } if (trimmed_primary_keys.empty()) { @@ -170,16 +173,22 @@ Result> RealtimePrimaryKeyWriter::Crea DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; prepared_fields.insert(prepared_fields.end(), write_schema->fields().begin(), write_schema->fields().end()); + const RealtimePartitionBucket partition_bucket(partition, bucket); + const int64_t initial_max_sequence_number = + realtime_context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + restored_max_sequence_number); return std::shared_ptr(new RealtimePrimaryKeyWriter( - store_state.store, merge_tree_writer, write_schema, + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, arrow::schema(std::move(prepared_fields)), arrow::schema(std::move(key_fields)), trimmed_primary_keys, key_comparator, store_state.initial_offset, - restored_max_sequence_number, memory_pool)); + initial_max_sequence_number, memory_pool)); } RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( const std::shared_ptr& realtime_store, const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, const std::shared_ptr& write_schema, const std::shared_ptr& prepared_schema, const std::shared_ptr& key_schema, @@ -190,6 +199,8 @@ RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(realtime_store), merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), write_schema_(write_schema), prepared_schema_(prepared_schema), key_schema_(key_schema), @@ -237,6 +248,8 @@ Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ + count)})); next_offset_ += count; last_sequence_number_ += count; + realtime_context_->AdvanceMaterializedMaxSequenceNumber(partition_bucket_, + last_sequence_number_); return Status::OK(); } diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index 6abb1ccd0..9a5aa4c68 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -26,6 +26,7 @@ #include #include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" namespace arrow { @@ -38,16 +39,19 @@ namespace paimon { class MemoryPool; class MergeTreeWriter; class FieldsComparator; +class RealtimeContextImpl; struct RealtimeStoreState; /// Coordinates framework-prepared primary-key real-time writes. class RealtimePrimaryKeyWriter final : public BatchWriter { public: static Result> Create( + const std::map& partition, int32_t bucket, const std::shared_ptr& write_schema, const std::vector& trimmed_primary_keys, const std::shared_ptr& key_comparator, - const RealtimeStoreState& store_state, int64_t restore_max_sequence_number, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, const std::shared_ptr& merge_tree_writer, const std::shared_ptr& memory_pool); @@ -64,6 +68,8 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { private: RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, const std::shared_ptr& write_schema, const std::shared_ptr& prepared_schema, const std::shared_ptr& key_schema, @@ -79,6 +85,8 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { std::shared_ptr arrow_pool_; std::shared_ptr realtime_store_; std::shared_ptr merge_tree_writer_; + std::shared_ptr realtime_context_; + RealtimePartitionBucket partition_bucket_; std::shared_ptr write_schema_; std::shared_ptr prepared_schema_; std::shared_ptr key_schema_; diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index c61b04b0d..e68ff670d 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1526,6 +1526,51 @@ TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector first_rows = { + {0, "value-0", "p0"}, {1, "value-1", "p0"}, {2, "value-2", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, first_progress.size()); + ASSERT_EQ(OffsetRange(0, 3), first_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(first_progress).size()); + ASSERT_EQ(0, NewFiles(first_progress)[0]->min_sequence_number); + ASSERT_EQ(2, NewFiles(first_progress)[0]->max_sequence_number); + ASSERT_OK(first_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector second_rows = {{0, "updated-0", "p0"}, {3, "value-3", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_progress.size()); + ASSERT_EQ(OffsetRange(3, 5), second_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(second_progress).size()); + ASSERT_EQ(3, NewFiles(second_progress)[0]->min_sequence_number); + ASSERT_EQ(4, NewFiles(second_progress)[0]->max_sequence_number); + + first_progress.push_back(std::move(second_progress[0])); + ASSERT_OK(Commit(first_progress, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{0, "updated-0", "p0"}, + {1, "value-1", "p0"}, + {2, "value-2", "p0"}, + {3, "value-3", "p0"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkPartitionBucketRecovery) { options_[Options::BUCKET] = "2"; CreatePkTable(/*partition_keys=*/{"pt"}); From a4f9a0c91e093f3130bd004f0134cf59f4c05697 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 15:56:00 +0800 Subject: [PATCH 19/24] refactor(realtime): simplify primary key merge readers --- .../core/operation/merge_file_split_read.cpp | 196 +++--------------- .../core/operation/merge_file_split_read.h | 9 +- .../table/source/key_value_table_read.cpp | 21 +- 3 files changed, 35 insertions(+), 191 deletions(-) diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index 2f64f6df8..c85e75ee0 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -30,7 +30,6 @@ #include "arrow/c/bridge.h" #include "arrow/type.h" #include "fmt/format.h" -#include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/table/special_fields.h" @@ -79,82 +78,36 @@ struct KeyValue; template class MergeFunctionWrapper; -namespace { - -class ConcatNonOverlappingMergeReaders final : public SortMergeReader { - public: - explicit ConcatNonOverlappingMergeReaders( - std::vector>&& readers) - : readers_(std::move(readers)) {} - - Result> NextBatch() override { - while (current_ < readers_.size()) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, - readers_[current_]->NextBatch()); - if (iterator) { - return iterator; - } - readers_[current_]->Close(); - ++current_; - } - return std::unique_ptr(); - } - - void Close() override { - while (current_ < readers_.size()) { - readers_[current_++]->Close(); - } - } - - std::shared_ptr GetReaderMetrics() const override { - return MetricsImpl::CollectReadMetrics(readers_); - } - - private: - std::vector> readers_; - size_t current_ = 0; -}; - -} - class MergeFileSplitRead::RealtimeReaderBuilder { public: static Result> Create( MergeFileSplitRead* owner, const std::vector>& disk_splits, - std::vector&& additional_readers) { + std::vector>&& additional_readers) { RealtimeReaderBuilder builder(owner); - if (disk_splits.empty()) { - std::vector> readers; - readers.reserve(additional_readers.size()); - for (AdditionalKeyValueReader& additional : additional_readers) { - readers.push_back(std::move(additional.reader)); - } - return builder.CreateMergedReader(std::move(readers)); + std::vector> readers; + if (!disk_splits.empty()) { + PAIMON_RETURN_NOT_OK(builder.CollectDiskReaders(disk_splits, &readers)); } - - PAIMON_RETURN_NOT_OK(builder.CollectDiskInputs(disk_splits)); - builder.AddRangeInputs(std::move(additional_readers)); - return builder.CreateReader(); + readers.reserve(readers.size() + additional_readers.size()); + for (std::unique_ptr& additional_reader : additional_readers) { + readers.push_back(std::move(additional_reader)); + } + return builder.CreateMergedReader(std::move(readers)); } private: - struct RangeInput { - std::shared_ptr min_key; - std::shared_ptr max_key; - std::vector disk_runs; - std::unique_ptr additional_reader; - }; - explicit RealtimeReaderBuilder(MergeFileSplitRead* owner) : owner_(owner) {} - Status CollectDiskInputs(const std::vector>& disk_splits) { - first_split_ = std::dynamic_pointer_cast(disk_splits.front()); - if (!first_split_) { + Status CollectDiskReaders(const std::vector>& disk_splits, + std::vector>* readers) { + std::shared_ptr first_split = + std::dynamic_pointer_cast(disk_splits.front()); + if (!first_split) { return Status::Invalid("merge input disk split is not a data split"); } - const BinaryRow& partition = first_split_->Partition(); - const int32_t bucket = first_split_->Bucket(); - PAIMON_ASSIGN_OR_RAISE(data_file_path_factory_, + const BinaryRow& partition = first_split->Partition(); + const int32_t bucket = first_split->Bucket(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); std::vector> data_files; @@ -187,46 +140,23 @@ class MergeFileSplitRead::RealtimeReaderBuilder { } } - dv_factory_ = DeletionVector::CreateFactory( + DeletionVector::Factory dv_factory = DeletionVector::CreateFactory( owner_->options_.GetFileSystem(), DeletionVector::CreateDeletionFileMap(data_files, deletion_files), owner_->pool_); std::vector> disk_sections = IntervalPartition(data_files, owner_->key_comparator_).Partition(); - inputs_.reserve(disk_sections.size()); - for (std::vector& section : disk_sections) { - std::shared_ptr min_file = section.front().Files().front(); - std::shared_ptr max_file = min_file; + for (const std::vector& section : disk_sections) { for (const SortedRun& run : section) { - for (const std::shared_ptr& file : run.Files()) { - if (owner_->key_comparator_->CompareTo(file->min_key, min_file->min_key) < 0) { - min_file = file; - } - if (owner_->key_comparator_->CompareTo(file->max_key, max_file->max_key) > 0) { - max_file = file; - } - } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + owner_->CreateReaderForRun(partition, run, dv_factory, + owner_->predicate_for_keys_, + data_file_path_factory)); + readers->push_back(std::move(disk_reader)); } - inputs_.push_back(RangeInput{std::shared_ptr(min_file, &min_file->min_key), - std::shared_ptr(max_file, &max_file->max_key), - std::move(section), nullptr}); } return Status::OK(); } - void AddRangeInputs(std::vector&& additional_readers) { - inputs_.reserve(inputs_.size() + additional_readers.size()); - for (AdditionalKeyValueReader& additional : additional_readers) { - has_unknown_range_ |= !additional.min_key || !additional.max_key; - inputs_.push_back(RangeInput{additional.min_key, additional.max_key, {}, - std::move(additional.reader)}); - } - } - - Result> CreateDiskReader(const SortedRun& run) { - return owner_->CreateReaderForRun(first_split_->Partition(), run, dv_factory_, - owner_->predicate_for_keys_, data_file_path_factory_); - } - Result> CreateMergedReader( std::vector>&& record_readers) { if (record_readers.empty()) { @@ -265,83 +195,7 @@ class MergeFileSplitRead::RealtimeReaderBuilder { owner_->pool_); } - Result> CreateUnknownRangeReader() { - std::vector> readers; - for (RangeInput& input : inputs_) { - for (const SortedRun& run : input.disk_runs) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, - CreateDiskReader(run)); - readers.push_back(std::move(disk_reader)); - } - if (input.additional_reader) { - readers.push_back(std::move(input.additional_reader)); - } - } - return CreateMergedReader(std::move(readers)); - } - - Result> CreateKnownRangeReader() { - std::sort(inputs_.begin(), inputs_.end(), - [this](const RangeInput& lhs, const RangeInput& rhs) { - return owner_->key_comparator_->CompareTo(*lhs.min_key, *rhs.min_key) < 0; - }); - std::vector> components; - std::shared_ptr component_max_key; - for (RangeInput& input : inputs_) { - if (components.empty() || - owner_->key_comparator_->CompareTo(*input.min_key, *component_max_key) > 0) { - components.emplace_back(); - component_max_key = input.max_key; - } else if (owner_->key_comparator_->CompareTo(*input.max_key, *component_max_key) > 0) { - component_max_key = input.max_key; - } - components.back().push_back(std::move(input)); - } - - std::vector> component_readers; - component_readers.reserve(components.size()); - for (std::vector& component : components) { - if (component.size() == 1 && !component.front().additional_reader) { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr disk_component, - owner_->CreateSortMergeReaderForSection( - component.front().disk_runs, first_split_->Partition(), dv_factory_, - component.front().disk_runs.size() == 1 ? owner_->context_->GetPredicate() - : owner_->predicate_for_keys_, - data_file_path_factory_, false)); - component_readers.push_back(std::move(disk_component)); - continue; - } - - std::vector> readers; - for (RangeInput& input : component) { - for (const SortedRun& run : input.disk_runs) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, - CreateDiskReader(run)); - readers.push_back(std::move(disk_reader)); - } - if (input.additional_reader) { - readers.push_back(std::move(input.additional_reader)); - } - } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr component_reader, - owner_->CreateSortMergeReader(std::move(readers))); - component_readers.push_back(std::move(component_reader)); - } - return CreateProjectedReader( - std::make_unique(std::move(component_readers))); - } - - Result> CreateReader() { - return has_unknown_range_ ? CreateUnknownRangeReader() : CreateKnownRangeReader(); - } - MergeFileSplitRead* owner_; - std::shared_ptr first_split_; - std::shared_ptr data_file_path_factory_; - DeletionVector::Factory dv_factory_; - std::vector inputs_; - bool has_unknown_range_ = false; }; Result> MergeFileSplitRead::Create( @@ -426,7 +280,7 @@ Result> MergeFileSplitRead::CreateReader( Result> MergeFileSplitRead::CreateRealtimeReader( const std::vector>& disk_splits, - std::vector&& additional_readers) { + std::vector>&& additional_readers) { return RealtimeReaderBuilder::Create(this, disk_splits, std::move(additional_readers)); } diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 6c1399978..3cc63a444 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -55,7 +55,6 @@ class FieldsComparator; class FileBatchReader; class FileStorePathFactory; class InternalReadContext; -class InternalRow; class MemoryPool; class SchemaManager; class SortedRun; @@ -66,12 +65,6 @@ struct KeyValue; template class MergeFunctionWrapper; -struct AdditionalKeyValueReader { - std::unique_ptr reader; - std::shared_ptr min_key; - std::shared_ptr max_key; -}; - /// If the class name below is enclosed in parentheses, it might be present in the read path; /// otherwise, it must be present in the read path. /// @@ -130,7 +123,7 @@ class MergeFileSplitRead : public AbstractSplitRead { Result> CreateRealtimeReader( const std::vector>& disk_splits, - std::vector&& additional_readers); + std::vector>&& additional_readers); void SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper); diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 31779e049..dc69ebb55 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -55,7 +55,7 @@ struct ColumnarBatchContext; namespace { -Result> CreateMemoryReaders( +Result>> CreateMemoryReaders( const std::shared_ptr& split, const RealtimePartitionBucketView& memory, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, @@ -76,9 +76,8 @@ Result> CreateMemoryReaders( PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema, c_schema.get())); ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; - PAIMON_ASSIGN_OR_RAISE( - std::vector> batch_readers, - memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, + memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); ScopeGuard batch_readers_guard([&batch_readers]() { for (const std::unique_ptr& reader : batch_readers) { if (reader) { @@ -86,7 +85,7 @@ Result> CreateMemoryReaders( } } }); - std::vector result; + std::vector> result; result.reserve(batch_readers.size()); for (std::unique_ptr& reader : batch_readers) { if (!reader) { @@ -98,17 +97,15 @@ Result> CreateMemoryReaders( split->MemoryEndOffset()), key_schema, value_schema, memory_pool)); auto merge = std::make_unique(false); - result.push_back(AdditionalKeyValueReader{ - std::make_unique( - std::move(prepared_reader), key_comparator, - std::make_shared(std::move(merge))), - nullptr, nullptr}); + result.push_back(std::make_unique( + std::move(prepared_reader), key_comparator, + std::make_shared(std::move(merge)))); } batch_readers_guard.Release(); return result; } -} +} // namespace KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, @@ -268,7 +265,7 @@ Result> KeyValueTableRead::CreateRealtimeReader( auto* merge_read = dynamic_cast(read.get()); if (merge_read) { PAIMON_ASSIGN_OR_RAISE( - std::vector memory_readers, + std::vector> memory_readers, CreateMemoryReaders(realtime_split, memory, merge_read->GetKeySchema(), merge_read->GetValueSchema(), merge_read->GetKeyComparator(), context_, GetMemoryPool())); From d983089057025ec3e85a09a338f90c1826db18ce Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:20:13 +0800 Subject: [PATCH 20/24] fix(realtime): validate PK reader contracts --- include/paimon/realtime/realtime_store.h | 11 +- .../merged_key_value_record_reader_test.cpp | 98 ++---- .../core/operation/file_store_write.cpp | 2 +- .../operation/key_value_file_store_write.cpp | 2 +- .../realtime/arrow_realtime_store_factory.cpp | 5 +- .../realtime/prepared_key_value_reader.cpp | 231 ++++++++++++-- .../core/realtime/prepared_key_value_reader.h | 22 +- .../realtime/primary_key_realtime_store.cpp | 160 ++++++++-- .../realtime/primary_key_realtime_store.h | 4 +- .../primary_key_realtime_store_test.cpp | 144 ++++++++- .../core/realtime/realtime_context_impl.cpp | 55 +++- .../core/realtime/realtime_context_impl.h | 12 +- .../core/realtime/realtime_context_test.cpp | 27 +- .../realtime/realtime_primary_key_writer.cpp | 32 +- .../realtime/realtime_primary_key_writer.h | 2 +- src/paimon/core/realtime/realtime_reader.h | 11 + .../core/realtime/realtime_reader_test.cpp | 27 +- .../table/source/key_value_table_read.cpp | 11 +- src/paimon/core/table/source/table_scan.cpp | 2 +- test/inte/realtime_write_inte_test.cpp | 300 +++++++++++++++++- 20 files changed, 971 insertions(+), 187 deletions(-) diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index dc5d543ac..792bb1c56 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -47,7 +47,10 @@ struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig { StatisticsMode statistics_mode; }; -struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig {}; +struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig { + /// Primary-key fields after removing partition fields, in comparison order. + std::vector trimmed_primary_keys; +}; using RealtimeStoreCreateConfig = std::variant; @@ -148,7 +151,8 @@ class PAIMON_EXPORT RealtimeStore { /// including across `NextBatch` boundaries, is sorted by full primary key then sequence /// number; all readers collectively cover sealed mutations exactly once. Reader cardinality is /// independent of the number of writes. Paimon adapts and merges those rows before writing - /// files. + /// files. Paimon validates the complete ordering and coverage before publishing generated file + /// state; a violation fails the prepare operation. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -168,7 +172,8 @@ class PAIMON_EXPORT RealtimeStore { /// contain multiple mutations per key. Each returned primary-key reader's complete stream is /// sorted by full primary key then sequence number, and all readers collectively cover raw /// mutations exactly once. Reader cardinality is independent of the number of writes. Paimon - /// retains `view` for the lifetime of the resulting framework reader. + /// validates ordering while adapting each complete reader stream and retains `view` for the + /// lifetime of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 21b0a16b1..775f271e3 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -95,7 +95,7 @@ class TrackingBatchReader : public BatchReader { int32_t* close_count_; }; -} +} // namespace class MergedKeyValueRecordReaderTest : public testing::Test { public: @@ -229,8 +229,7 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderOffsetFilter) { ])") .ValueOrDie()); - auto batch_reader = - std::make_unique(prepared_array, prepared_type, 2); + auto batch_reader = std::make_unique(prepared_array, prepared_type, 2); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), @@ -248,77 +247,28 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderOffsetFilter) { KeyValueChecker::CheckResult(expected, results, 1, 2); } -TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderRawCountBeforeDedup) { - std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), - DataField(1, arrow::field("v0", arrow::int32()))}; +TEST_F(MergedKeyValueRecordReaderTest, TestRejectsUnsortedPluginRowsAcrossBatches) { + std::vector value_fields = {DataField(0, arrow::field("id", arrow::int32()))}; std::shared_ptr value_schema = DataField::ConvertDataFieldsToArrowSchema(value_fields); std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); - auto prepared_array = std::dynamic_pointer_cast( - arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ - [0, 10, 0, 1, 100], - [2, 11, 1, 1, 101], - [0, 12, 2, 2, 200] - ])") - .ValueOrDie()); - - int64_t raw_row_count = 0; - auto batch_reader = - std::make_unique(prepared_array, prepared_type, 2); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr raw_reader, - AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, std::nullopt, key_schema, - value_schema, pool_, &raw_row_count)); - - ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, - FieldsComparator::Create({value_fields[0]}, true)); - auto merged_reader = std::make_unique( - std::move(raw_reader), key_comparator, merge_function_wrapper_); - ASSERT_OK_AND_ASSIGN( - std::vector results, - (ReadResultCollector::CollectKeyValueResult< - MergedKeyValueRecordReader, KeyValueRecordReader::Iterator>(merged_reader.get()))); - - ASSERT_EQ(raw_row_count, 3); - std::vector row_kinds = {const_cast(RowKind::UpdateAfter()), - const_cast(RowKind::Insert())}; - std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; - std::vector expected = KeyValueChecker::GenerateKeyValues( - row_kinds, {11, 12}, levels, {{1}, {2}}, {{1, 101}, {2, 200}}, pool_); - KeyValueChecker::CheckResult(expected, results, 1, 2); -} - -TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderRawCountBeforeFilter) { - std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32()))}; - std::shared_ptr value_schema = - DataField::ConvertDataFieldsToArrowSchema(value_fields); - std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); - std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); - auto prepared_array = std::dynamic_pointer_cast( + std::shared_ptr prepared_array = arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ - [0, 10, 0, 1], - [0, 11, 1, 2], - [0, 12, 2, 3], - [0, 13, 3, 4] - ])") - .ValueOrDie()); - - int64_t raw_row_count = 0; + [0, 10, 0, 2], + [0, 11, 1, 1] + ])") + .ValueOrDie(); auto batch_reader = - std::make_unique(prepared_array, prepared_type, 2); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr reader, - AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), - value_schema, value_schema, pool_, &raw_row_count)); - ASSERT_OK_AND_ASSIGN( - std::vector results, - (ReadResultCollector::CollectKeyValueResult(reader.get()))); - - ASSERT_EQ(results.size(), 2); - ASSERT_EQ(raw_row_count, 4); + std::make_unique(prepared_array, prepared_type, /*batch_size=*/1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, + std::nullopt, key_schema, value_schema, pool_)); + Result> result = + ReadResultCollector::CollectKeyValueResult(reader.get()); + ASSERT_NOK_WITH_MSG(result, "not globally sorted by primary key and sequence number"); } TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { @@ -345,8 +295,7 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { ASSERT_EQ(query_results[0].value->GetFieldCount(), 1); ASSERT_EQ(query_results[0].value->GetInt(0), 1); - auto batch_reader = - std::make_unique(prepared_array, prepared_type, 1); + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); ASSERT_NOK_WITH_MSG(AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, std::nullopt, value_schema, value_schema, pool_), "exact"); @@ -364,8 +313,7 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderSafeDecode) { auto invalid_array = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[0, 10, 0, 1]])").ValueOrDie()); - auto batch_reader = - std::make_unique(invalid_array, invalid_type, 1); + auto batch_reader = std::make_unique(invalid_array, invalid_type, 1); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), @@ -398,9 +346,12 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); auto prepared_array = std::dynamic_pointer_cast( arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ - [0, 10, 0, 1, [[100, 200], [300, 400]], [["k1", [7, 8]], ["k2", [9, 10]]], [[[11, 12], 13], [[21, 22], 23]]] + [0, 9, 9, 0, [[1, 2]], [["prefix", [3, 4]]], [[[5, 6], 7]]], + [0, 10, 0, 1, [[100, 200], [300, 400]], [["k1", [7, 8]], ["k2", [9, 10]]], [[[11, 12], 13], [[21, 22], 23]]], + [0, 11, 11, 2, [[8, 9]], [["suffix", [10, 11]]], [[[12, 13], 14]]] ])") .ValueOrDie()); + prepared_array = checked_pointer_cast(prepared_array->Slice(1, 1)); std::shared_ptr query_item_b = MakeField("renamed_b", arrow::int32(), 11); std::shared_ptr query_item_a = MakeField("renamed_a", arrow::int32(), 10); @@ -420,8 +371,7 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { std::shared_ptr query_value_schema = arrow::schema({id, query_items, query_attrs, query_keyed_values}); - auto batch_reader = - std::make_unique(prepared_array, prepared_type, 1); + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 4d4f45156..84a324762 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -198,7 +198,7 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options)); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options, *schema)); if (ignore_previous_files) { return Status::NotImplemented( "PK realtime v1 requires restore from the latest snapshot"); 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 de7217ec3..ee6445057 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -155,7 +155,7 @@ Result> KeyValueFileStoreWrite::CreateWriter( RealtimeStoreState store_state, realtime_context_impl->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ std::move(c_write_schema), options_.ToMap(), pool_, partition_map, bucket, - PrimaryKeyRealtimeStoreCreateConfig{}})); + PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys}})); realtime_store_state = std::move(store_state); compact_manager = std::make_shared(); } else { diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 4cfdb4c3d..babc55a3d 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -50,8 +50,11 @@ Result> ArrowRealtimeStoreFactory::Create( request.memory_pool, arrow_pool); } + const PrimaryKeyRealtimeStoreCreateConfig& config = + std::get(request.mode_config); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool)); + PrimaryKeyRealtimeStore::Create( + imported_schema, config.trimmed_primary_keys, request.memory_pool)); return std::shared_ptr(std::move(store)); } diff --git a/src/paimon/core/realtime/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp index b99f67dd9..6b3afcd19 100644 --- a/src/paimon/core/realtime/prepared_key_value_reader.cpp +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -29,9 +30,11 @@ #include "arrow/array/array_nested.h" #include "arrow/array/array_primitive.h" #include "arrow/array/builder_primitive.h" +#include "arrow/buffer.h" #include "arrow/c/bridge.h" #include "arrow/compute/api.h" #include "arrow/type.h" +#include "arrow/util/bit_util.h" #include "fmt/format.h" #include "paimon/common/data/columnar/columnar_batch_context.h" #include "paimon/common/data/columnar/columnar_row_ref.h" @@ -42,6 +45,7 @@ #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/realtime/realtime_fields.h" #include "paimon/core/utils/nested_projection_utils.h" @@ -61,6 +65,68 @@ constexpr int32_t kPreparedValueStartIndex = 3; Result> AlignArrayByPaimonIds( const std::shared_ptr& array, const std::shared_ptr& read_type); +class RealtimeOffsetCoverage { + public: + static Result> Create( + const OffsetRange& sealed_offsets, size_t reader_count, + const std::shared_ptr& arrow_pool) { + if (sealed_offsets.begin < 0 || sealed_offsets.end < sealed_offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr seen_offsets, + arrow::AllocateEmptyBitmap(sealed_offsets.Count(), arrow_pool.get())); + return std::shared_ptr(new RealtimeOffsetCoverage( + sealed_offsets, reader_count, std::move(seen_offsets), arrow_pool)); + } + + Status Add(const arrow::Int64Array& offsets) { + std::lock_guard lock(mutex_); + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (offset < sealed_offsets_.begin || offset >= sealed_offsets_.end) { + return Status::Invalid( + "PK real-time store commit reader offset is outside the sealed range"); + } + const int64_t index = offset - sealed_offsets_.begin; + if (arrow::bit_util::GetBit(seen_offsets_->data(), index)) { + return Status::Invalid( + "PK real-time store commit readers contain duplicate REALTIME_OFFSET"); + } + arrow::bit_util::SetBit(seen_offsets_->mutable_data(), index); + ++seen_count_; + } + return Status::OK(); + } + + Status FinishReader() { + std::lock_guard lock(mutex_); + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && seen_count_ != sealed_offsets_.Count()) { + return Status::Invalid( + "PK real-time store commit readers did not cover the sealed range"); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t reader_count, + std::shared_ptr seen_offsets, + const std::shared_ptr& arrow_pool) + : sealed_offsets_(sealed_offsets), + reader_count_(reader_count), + arrow_pool_(arrow_pool), + seen_offsets_(std::move(seen_offsets)) {} + + OffsetRange sealed_offsets_; + size_t reader_count_; + std::shared_ptr arrow_pool_; + std::shared_ptr seen_offsets_; + int64_t seen_count_ = 0; + size_t finished_reader_count_ = 0; + std::mutex mutex_; +}; + Status CheckPreparedField(const std::shared_ptr& schema, int32_t field_idx, const DataField& expected_field) { if (schema->num_fields() <= field_idx) { @@ -210,16 +276,20 @@ Result> AlignStructArrayByPaimonIds( return Status::Invalid( fmt::format("cannot find field id {} in prepared value struct", read_field_id)); } - std::shared_ptr child = array->field(data_iter->second); + std::shared_ptr child = + arrow::MakeArray(array->data()->child_data[data_iter->second]); PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type())); aligned_arrays.push_back(std::move(child)); } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr aligned, - arrow::StructArray::Make(aligned_arrays, read_type->fields(), array->null_bitmap(), - array->null_count(), array->offset())); - return aligned; + std::shared_ptr aligned_data = array->data()->Copy(); + aligned_data->type = read_type; + aligned_data->child_data.clear(); + aligned_data->child_data.reserve(aligned_arrays.size()); + for (const std::shared_ptr& aligned_array : aligned_arrays) { + aligned_data->child_data.push_back(aligned_array->data()); + } + return arrow::MakeArray(std::move(aligned_data)); } Result> AlignListArrayByPaimonIds( @@ -351,15 +421,18 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { const std::optional& visible_offsets, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, - const std::shared_ptr& pool, int64_t* raw_row_count) + const std::shared_ptr& key_comparator, + const std::shared_ptr& pool, + const std::shared_ptr& offset_coverage) : reader_(std::move(reader)), prepared_schema_(prepared_schema), visible_offsets_(visible_offsets), key_schema_(key_schema), value_schema_(value_schema), + key_comparator_(key_comparator), pool_(pool), arrow_pool_(GetArrowPool(pool)), - raw_row_count_(raw_row_count) {} + offset_coverage_(offset_coverage) {} ~PreparedKeyValueReader() override { Close(); @@ -425,6 +498,10 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { ResetBatchState(); PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); if (BatchReader::IsEofBatch(batch)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } return std::unique_ptr(); } auto& [c_array, c_schema] = batch; @@ -436,17 +513,14 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { std::shared_ptr data_batch = checked_pointer_cast(arrow_array); PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch)); - if (raw_row_count_ != nullptr) { - int64_t updated_count = 0; - if (__builtin_add_overflow(*raw_row_count_, data_batch->length(), &updated_count)) { - return Status::Invalid("prepared raw row count overflow"); - } - *raw_row_count_ = updated_count; - } + PAIMON_RETURN_NOT_OK(ValidateOrdering(data_batch)); std::shared_ptr> offset_array = checked_pointer_cast>( data_batch->field(kRealtimeOffsetIndex)); + if (offset_coverage_) { + PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array)); + } PAIMON_ASSIGN_OR_RAISE( data_batch, ApplyOffsetFilter(data_batch, offset_array, visible_offsets_, arrow_pool_.get())); @@ -504,6 +578,36 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { return Status::OK(); } + Status ValidateOrdering(const std::shared_ptr& data_batch) { + if (data_batch->length() == 0) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_)); + std::shared_ptr key_context = + std::make_shared(key_fields, pool_); + std::shared_ptr sequences = + checked_pointer_cast(data_batch->field(kSequenceNumberIndex)); + for (int64_t row = 0; row < data_batch->length(); ++row) { + ColumnarRowRef current_key(key_context, row); + if (previous_key_context_) { + ColumnarRowRef previous_key(previous_key_context_, previous_key_row_); + const int32_t key_comparison = + key_comparator_->CompareTo(previous_key, current_key); + if (key_comparison > 0 || + (key_comparison == 0 && previous_sequence_ > sequences->Value(row))) { + return Status::Invalid( + "PK real-time plugin reader is not globally sorted by primary key and " + "sequence number"); + } + } + previous_key_context_ = key_context; + previous_key_row_ = row; + previous_sequence_ = sequences->Value(row); + } + return Status::OK(); + } + void ResetBatchState() { key_ctx_.reset(); value_ctx_.reset(); @@ -518,23 +622,32 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { std::optional visible_offsets_; std::shared_ptr key_schema_; std::shared_ptr value_schema_; + std::shared_ptr key_comparator_; std::shared_ptr pool_; std::shared_ptr arrow_pool_; - int64_t* raw_row_count_; + std::shared_ptr offset_coverage_; + bool offset_coverage_finished_ = false; std::shared_ptr key_ctx_; std::shared_ptr value_ctx_; std::shared_ptr> row_kind_array_; std::shared_ptr> sequence_number_array_; + std::shared_ptr previous_key_context_; + int64_t previous_key_row_ = 0; + int64_t previous_sequence_ = 0; }; -} +} // namespace -Result> AdaptPreparedBatchReader( +namespace { + +Result> AdaptPreparedBatchReaderImpl( std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, const std::optional& visible_offsets, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, - const std::shared_ptr& memory_pool, int64_t* raw_row_count) { + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool, + const std::shared_ptr& offset_coverage) { std::unique_ptr owned_reader = std::move(reader); if (!owned_reader) { return Status::Invalid("prepared batch reader cannot be null"); @@ -547,6 +660,9 @@ Result> AdaptPreparedBatchReader( if (!value_schema) { return Status::Invalid("prepared value schema cannot be null"); } + if (!key_comparator) { + return Status::Invalid("prepared key comparator cannot be null"); + } if (!memory_pool) { return Status::Invalid("prepared reader memory pool cannot be null"); } @@ -555,11 +671,82 @@ Result> AdaptPreparedBatchReader( if (!visible_offsets.has_value()) { PAIMON_RETURN_NOT_OK(ValidateExactCommitSchema(prepared_schema, value_schema)); } - std::unique_ptr result( - new PreparedKeyValueReader(std::move(owned_reader), prepared_schema, visible_offsets, - key_schema, value_schema, memory_pool, raw_row_count)); + std::unique_ptr result(new PreparedKeyValueReader( + std::move(owned_reader), prepared_schema, visible_offsets, key_schema, value_schema, + key_comparator, memory_pool, offset_coverage)); close_guard.Release(); return result; } +} // namespace + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) { + return AdaptPreparedBatchReaderImpl(std::move(reader), prepared_schema, visible_offsets, + key_schema, value_schema, key_comparator, memory_pool, + /*offset_coverage=*/nullptr); +} + +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) { + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(sealed_offsets, readers.size(), arrow_pool)); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); + std::vector> adapted_readers; + ScopeGuard adapted_readers_guard([&adapted_readers]() { + for (const std::unique_ptr& reader : adapted_readers) { + reader->Close(); + } + }); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr adapted_reader, + AdaptPreparedBatchReaderImpl( + std::move(reader), prepared_schema, std::nullopt, key_schema, + value_schema, key_comparator, memory_pool, offset_coverage)); + adapted_readers.push_back(std::move(adapted_reader)); + } + readers_guard.Release(); + adapted_readers_guard.Release(); + return adapted_readers; +} + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + if (!key_schema) { + return Status::Invalid("prepared key schema cannot be null"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector key_fields, + DataField::ConvertArrowSchemaToDataFields(key_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); + return AdaptPreparedBatchReader(std::move(reader), prepared_schema, visible_offsets, key_schema, + value_schema, key_comparator, memory_pool); } + +} // namespace paimon diff --git a/src/paimon/core/realtime/prepared_key_value_reader.h b/src/paimon/core/realtime/prepared_key_value_reader.h index e7a6f9651..064a62958 100644 --- a/src/paimon/core/realtime/prepared_key_value_reader.h +++ b/src/paimon/core/realtime/prepared_key_value_reader.h @@ -21,6 +21,7 @@ #include #include #include +#include #include "arrow/type_fwd.h" #include "paimon/core/io/key_value_record_reader.h" @@ -29,6 +30,7 @@ namespace paimon { class BatchReader; +class FieldsComparator; class MemoryPool; Result> AdaptPreparedBatchReader( @@ -36,6 +38,22 @@ Result> AdaptPreparedBatchReader( const std::optional& visible_offsets, const std::shared_ptr& key_schema, const std::shared_ptr& value_schema, - const std::shared_ptr& memory_pool, int64_t* raw_row_count = nullptr); + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool); -} +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool); + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 0d6de9f5f..f43f60472 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -18,20 +18,31 @@ #include "paimon/core/realtime/primary_key_realtime_store.h" +#include #include +#include #include +#include #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/types/data_field.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/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/macros.h" namespace paimon { -Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema) { if (options.GetBucket() <= 0) { return Status::NotImplemented("PK realtime v1 requires fixed buckets"); } @@ -60,6 +71,21 @@ Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) { options.GetChangelogProducer() != ChangelogProducer::NONE) { return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); } + PAIMON_ASSIGN_OR_RAISE(std::vector primary_key_fields, + schema.TrimmedPrimaryKeyFields()); + for (const DataField& field : primary_key_fields) { + if (field.Type()->id() == arrow::Type::FLOAT || field.Type()->id() == arrow::Type::DOUBLE) { + return Status::NotImplemented( + "PK realtime v1 does not support FLOAT or DOUBLE primary keys"); + } + } + if (options.GlobalIndexEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(schema)); + if (!definitions.Definitions().empty()) { + return Status::NotImplemented("PK realtime v1 does not support global indexes"); + } + } return Status::OK(); } @@ -128,14 +154,56 @@ class ReadView final : public RealtimeReadView { class RawBatchReader final : public BatchReader { public: - RawBatchReader(std::vector batches) - : batches_(std::move(batches)), metrics_(std::make_shared()) {} + RawBatchReader(std::vector batches, std::vector key_field_indexes, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) + : batches_(std::move(batches)), + positions_(batches_.size(), 0), + key_field_indexes_(std::move(key_field_indexes)), + key_comparator_(key_comparator), + memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + metrics_(std::make_shared()) { + key_contexts_.reserve(batches_.size()); + for (const StoredBatch& batch : batches_) { + arrow::ArrayVector key_arrays; + key_arrays.reserve(key_field_indexes_.size()); + for (int32_t field_index : key_field_indexes_) { + key_arrays.push_back(batch.data->field(field_index)); + } + key_contexts_.push_back( + std::make_shared(key_arrays, memory_pool_)); + } + } Result NextBatch() override { - if (next_ == batches_.size()) { + if (closed_) { + return MakeEofBatch(); + } + std::optional selected; + for (size_t i = 0; i < batches_.size(); ++i) { + if (positions_[i] >= batches_[i].data->length()) { + continue; + } + if (!selected.has_value() || Less(i, selected.value())) { + selected = i; + } + } + if (!selected.has_value()) { return MakeEofBatch(); } - const std::shared_ptr& batch = batches_[next_++].data; + const size_t batch_index = selected.value(); + arrow::Int64Builder index_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Append(positions_[batch_index])); + std::shared_ptr index; + PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Finish(&index)); + arrow::compute::ExecContext context(arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum taken, + arrow::compute::Take(arrow::Datum(batches_[batch_index].data), arrow::Datum(index), + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + std::shared_ptr batch = taken.make_array(); + ++positions_[batch_index]; auto array = std::make_unique(); auto schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, array.get(), schema.get())); @@ -146,12 +214,38 @@ class RawBatchReader final : public BatchReader { return metrics_; } void Close() override { + if (closed_) { + return; + } + closed_ = true; batches_.clear(); + positions_.clear(); + key_contexts_.clear(); } private: + bool Less(size_t left, size_t right) const { + ColumnarRowRef left_key(key_contexts_[left], positions_[left]); + ColumnarRowRef right_key(key_contexts_[right], positions_[right]); + const int32_t key_comparison = key_comparator_->CompareTo(left_key, right_key); + if (key_comparison != 0) { + return key_comparison < 0; + } + const std::shared_ptr left_sequences = + checked_pointer_cast(batches_[left].data->field(1)); + const std::shared_ptr right_sequences = + checked_pointer_cast(batches_[right].data->field(1)); + return left_sequences->Value(positions_[left]) < right_sequences->Value(positions_[right]); + } + + bool closed_ = false; std::vector batches_; - size_t next_ = 0; + std::vector positions_; + std::vector key_field_indexes_; + std::shared_ptr key_comparator_; + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::vector> key_contexts_; std::shared_ptr metrics_; }; @@ -159,8 +253,13 @@ class RawBatchReader final : public BatchReader { class PrimaryKeyRealtimeStore::Impl { public: - explicit Impl(std::shared_ptr prepared_schema) - : prepared_schema_(std::move(prepared_schema)) {} + Impl(std::shared_ptr prepared_schema, std::vector key_field_indexes, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) + : prepared_schema_(std::move(prepared_schema)), + key_field_indexes_(std::move(key_field_indexes)), + key_comparator_(key_comparator), + memory_pool_(memory_pool) {} Status Write(RealtimeWriteBatch&& write_batch) { if (!write_batch.batch || !write_batch.batch->GetData()) { @@ -212,9 +311,9 @@ class PrimaryKeyRealtimeStore::Impl { return Status::Invalid("segment was not created by the PK real-time store"); } std::vector> readers; - readers.reserve(segment->Batches().size()); - for (const StoredBatch& batch : segment->Batches()) { - readers.push_back(std::make_unique(std::vector{batch})); + if (!segment->Batches().empty()) { + readers.push_back(std::make_unique( + segment->Batches(), key_field_indexes_, key_comparator_, memory_pool_)); } return readers; } @@ -238,16 +337,13 @@ class PrimaryKeyRealtimeStore::Impl { return Status::Invalid("read view was not created by the PK real-time store"); } std::vector> readers; - size_t batch_count = 0; + std::vector batches; for (const std::shared_ptr& segment : typed->Segments()) { - batch_count += segment->Batches().size(); + batches.insert(batches.end(), segment->Batches().begin(), segment->Batches().end()); } - readers.reserve(batch_count); - for (const std::shared_ptr& segment : typed->Segments()) { - for (const StoredBatch& batch : segment->Batches()) { - readers.push_back( - std::make_unique(std::vector{batch})); - } + if (!batches.empty()) { + readers.push_back(std::make_unique( + std::move(batches), key_field_indexes_, key_comparator_, memory_pool_)); } return readers; } @@ -273,6 +369,9 @@ class PrimaryKeyRealtimeStore::Impl { private: std::shared_ptr prepared_schema_; + std::vector key_field_indexes_; + std::shared_ptr key_comparator_; + std::shared_ptr memory_pool_; mutable std::mutex mutex_; std::vector building_; std::vector> sealed_; @@ -286,12 +385,29 @@ PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; Result> PrimaryKeyRealtimeStore::Create( const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, const std::shared_ptr& memory_pool) { - if (!prepared_schema || !memory_pool) { + if (!prepared_schema || trimmed_primary_keys.empty() || !memory_pool) { return Status::Invalid("PK prepared schema or memory pool is null"); } - return std::shared_ptr( - new PrimaryKeyRealtimeStore(std::make_unique(prepared_schema))); + std::vector key_field_indexes; + std::vector key_fields; + key_field_indexes.reserve(trimmed_primary_keys.size()); + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + const int32_t field_index = prepared_schema->GetFieldIndex(key); + if (field_index < 3) { + return Status::Invalid("PK field is missing from prepared schema: ", key); + } + key_field_indexes.push_back(field_index); + PAIMON_ASSIGN_OR_RAISE(DataField field, DataField::ConvertArrowFieldToDataField( + prepared_schema->field(field_index))); + key_fields.push_back(std::move(field)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); + return std::shared_ptr(new PrimaryKeyRealtimeStore( + std::make_unique(prepared_schema, key_field_indexes, key_comparator, memory_pool))); } Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { return impl_->Write(std::move(batch)); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h index 5e18dd74f..d6a23ccf9 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.h +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -31,14 +31,16 @@ namespace paimon { class CoreOptions; class MemoryPool; +class TableSchema; -Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options); +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema); /// In-memory store for prepared primary-key real-time batches. class PrimaryKeyRealtimeStore final : public RealtimeStore { public: static Result> Create( const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, const std::shared_ptr& memory_pool); ~PrimaryKeyRealtimeStore() override; diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index 43831d7be..cafe3682e 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -31,6 +31,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/macros.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" @@ -44,7 +45,33 @@ std::shared_ptr PreparedSchema() { DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) ->WithNullable(false), DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false), - arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}); +} + +std::shared_ptr NestedPreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField(DataField( + 1, + arrow::field("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}))))}); +} + +std::shared_ptr PkSchema( + const std::shared_ptr& key_type = arrow::int64(), + const std::map& options = {}) { + return TableSchema::Create( + /*schema_id=*/0, + arrow::schema({arrow::field("id", key_type), arrow::field("value", arrow::utf8())}), + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options) + .value(); } std::unique_ptr MakeBatch(const std::string& json) { @@ -56,6 +83,27 @@ std::unique_ptr MakeBatch(const std::string& json) { return RecordBatchBuilder(c_array.get()).Finish().value(); } +std::unique_ptr MakeBatch(const std::shared_ptr& schema, + const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +void AssertOffsetsZero(const ArrowArray* array) { + ASSERT_NE(nullptr, array); + ASSERT_EQ(0, array->offset); + for (int64_t child = 0; child < array->n_children; ++child) { + AssertOffsetsZero(array->children[child]); + } + if (array->dictionary) { + AssertOffsetsZero(array->dictionary); + } +} + Result ReadJson(const std::vector>& readers) { std::vector> batches; for (const std::unique_ptr& reader : readers) { @@ -77,7 +125,7 @@ Result ReadJson(const std::vector>& re TEST(PrimaryKeyRealtimeStoreOptionsTest, TestSupportedOptions) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); - ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options)); + ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); } TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { @@ -95,13 +143,31 @@ TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { }; for (const std::map& option_map : unsupported_options) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); - ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options)); + ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); } } +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsFloatingPrimaryKeys) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float32())), + "FLOAT or DOUBLE primary keys"); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float64())), + "FLOAT or DOUBLE primary keys"); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsEnabledGlobalIndex) { + const std::map option_map = {{Options::BUCKET, "1"}, + {Options::PK_BTREE_INDEX_COLUMNS, "id"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK_WITH_MSG( + ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::int64(), option_map)), + "does not support global indexes"); +} + TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::optional> segment, store->SealForCommit()); ASSERT_FALSE(segment.has_value()); @@ -131,10 +197,11 @@ TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { } TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); ASSERT_OK(store->Write(RealtimeWriteBatch{ - MakeBatch(R"([[0, 5, 0, 3, "three"], [1, 6, 1, 1, "before"]])"), OffsetRange(0, 2)})); + MakeBatch(R"([[1, 6, 1, 1, "before"], [0, 5, 0, 3, "three"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write( RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, @@ -142,18 +209,45 @@ TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { ASSERT_TRUE(segment.has_value()); ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); ASSERT_EQ( - "-- is_valid: all not null\n-- child 0 type: int8\n [\n 0,\n 1,\n 2\n ]\n-- " - "child 1 type: int64\n [\n 5,\n 6,\n 7\n ]\n-- child 2 type: int64\n [\n " - "0,\n 1,\n 2\n ]\n-- child 3 type: int64\n [\n 3,\n 1,\n 2\n ]\n-- child " - "4 type: string\n [\n \"three\",\n \"before\",\n \"after\"\n ]", + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 2,\n 0\n ]\n-- " + "child 1 type: int64\n [\n 6,\n 7,\n 5\n ]\n-- child 2 type: int64\n [\n " + "1,\n 2,\n 0\n ]\n-- child 3 type: int64\n [\n 1,\n 2,\n 3\n ]\n-- child " + "4 type: string\n [\n \"before\",\n \"after\",\n \"three\"\n ]", actual); + readers[0]->Close(); + readers[0]->Close(); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, readers[0]->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); } -TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderExportsZeroOffsets) { + std::shared_ptr schema = NestedPreparedSchema(); ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(PreparedSchema(), GetDefaultPool())); + PrimaryKeyRealtimeStore::Create(schema, {"id"}, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(schema, R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]]])"), + OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + for (int32_t row = 0; row < 2; ++row) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + AssertOffsetsZero(batch.first.get()); + ASSERT_TRUE(arrow::ImportArray(batch.first.get(), batch.second.get()).ok()); + } +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); ASSERT_OK( store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, @@ -163,5 +257,27 @@ TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { ASSERT_EQ(std::optional(OffsetRange(4, 5)), view->GetOffsetRange()); } +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderCardinalityIsConstant) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 2, 1, 1, "one"]])"), OffsetRange(1, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + RealtimeQueryContext context{/*read_schema=*/nullptr, /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 415052a69..215e066ee 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -34,14 +34,33 @@ #include #include +#include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/arrow/abi.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/uuid.h" #include "paimon/macros.h" #include "paimon/realtime/realtime_store.h" #include "paimon/status.h" namespace paimon { +namespace { + +bool SameMode(const RealtimeStoreCreateConfig& left, const RealtimeStoreCreateConfig& right) { + if (left.index() != right.index()) { + return false; + } + if (const auto* left_pk = std::get_if(&left)) { + const auto& right_pk = std::get(right); + return left_pk->trimmed_primary_keys == right_pk.trimmed_primary_keys; + } + return true; +} + +} // namespace + RealtimeContextImpl::RealtimeContextImpl(const std::shared_ptr& factory) : factory_(factory) {} @@ -78,6 +97,14 @@ Status RealtimeContextImpl::Start() { Result RealtimeContextImpl::GetOrCreateRealtimeStore( RealtimeStoreCreateRequest&& request) { + if (!request.write_schema || !request.write_schema->release) { + return Status::Invalid("real-time store write schema is null"); + } + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr requested_schema, + arrow::ImportSchema(request.write_schema.get())); + schema_guard.Release(); std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); const RealtimePartitionBucket key(request.partition, request.bucket); @@ -86,19 +113,18 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { - if (request.write_schema) { - ArrowSchemaRelease(request.write_schema.get()); - } return Status::Invalid("real-time offset has reached INT64_MAX"); } initial_offset = offset_iter->second; } if (iter != stores_.end()) { - if (request.write_schema) { - ArrowSchemaRelease(request.write_schema.get()); + if (!SameMode(iter->second.mode_config, request.mode_config) || + !iter->second.write_schema->Equals(*requested_schema, /*check_metadata=*/true)) { + return Status::Invalid( + "real-time store schema or mode does not match the registered store"); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - iter->second->AcquireReadView()); + iter->second.store->AcquireReadView()); if (!read_view) { return Status::Invalid("real-time store returned a null read view"); } @@ -113,17 +139,18 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset}; + return RealtimeStoreState{iter->second.store, initial_offset}; } if (!request.memory_pool) { - if (request.write_schema) { - ArrowSchemaRelease(request.write_schema.get()); - } return Status::Invalid("real-time store memory pool is null"); } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*requested_schema, request.write_schema.get())); + RealtimeStoreCreateConfig mode_config = request.mode_config; Result> store_result = factory_->Create(std::move(request)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); - stores_.emplace(key, store); + stores_.emplace(key, + RealtimeStoreRegistryEntry{store, requested_schema, std::move(mode_config)}); if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); } @@ -147,9 +174,9 @@ Result> RealtimeContextImpl::AcquireRea result.reserve(stores_.size()); for (const auto& [partition_bucket, store] : stores_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - store->AcquireReadView()); + store.store->AcquireReadView()); result.push_back( - RealtimePartitionBucketView{partition_bucket, store, std::move(read_view)}); + RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; } @@ -266,7 +293,7 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, } auto store_iter = stores_.find(partition_bucket); if (store_iter != stores_.end()) { - notifications.emplace_back(partition_bucket, store_iter->second, + notifications.emplace_back(partition_bucket, store_iter->second.store, committed_end_offset); } } diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index aa4d263c6..9fa145e99 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -38,6 +38,10 @@ struct ArrowSchema; +namespace arrow { +class Schema; +} // namespace arrow + namespace paimon { class RealtimeStore; @@ -55,6 +59,12 @@ struct RealtimePartitionBucketView { std::shared_ptr read_view; }; +struct RealtimeStoreRegistryEntry { + std::shared_ptr store; + std::shared_ptr write_schema; + RealtimeStoreCreateConfig mode_config; +}; + class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { public: static Result> Create( @@ -98,7 +108,7 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::shared_ptr factory_; std::mutex mutex_; std::mutex progress_mutex_; - std::map> stores_; + std::map stores_; std::map materialized_max_sequence_numbers_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 15066ca1d..2b47e9dc9 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -96,10 +96,12 @@ class TestingRealtimeStoreFactory : public RealtimeStoreFactory { std::vector> stores; }; -std::unique_ptr MakeWriteSchema() { +std::unique_ptr MakeWriteSchema( + const std::shared_ptr& id_type = arrow::int64(), + const std::shared_ptr& metadata = nullptr) { auto schema = std::make_unique(); EXPECT_TRUE( - arrow::ExportSchema(*arrow::schema({arrow::field("id", arrow::int64())}), schema.get()) + arrow::ExportSchema(*arrow::schema({arrow::field("id", id_type)}, metadata), schema.get()) .ok()); return schema; } @@ -158,6 +160,27 @@ TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { ASSERT_EQ(1, factory->stores[2]->acquire_count); } +TEST(RealtimeContextTest, TestRejectsMismatchedSchemaOnStoreReuse) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + std::shared_ptr metadata = + arrow::key_value_metadata({"identity"}, {"v1"}); + ASSERT_OK(GetOrCreateAppendStore( + context, partition, 0, MakeWriteSchema(arrow::int64(), metadata), {}, GetDefaultPool())); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(arrow::int32(), metadata), {}, + GetDefaultPool()), + "schema or mode does not match"); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore( + context, partition, 0, + MakeWriteSchema(arrow::int64(), arrow::key_value_metadata({"identity"}, {"v2"})), {}, + GetDefaultPool()), + "schema or mode does not match"); + ASSERT_EQ(1, factory->stores.size()); +} + TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index b53831f0a..82318eadb 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -266,15 +266,12 @@ Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compac return Status::Invalid("PK real-time store sealed a null segment"); } std::optional sealed_range; - int64_t expected_raw_row_count = 0; if (segment) { sealed_range = segment.value()->GetOffsetRange(); - if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin || - __builtin_sub_overflow(sealed_range->end, sealed_range->begin, - &expected_raw_row_count)) { + if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin) { return Status::Invalid("PK real-time store returned an invalid sealed offset range"); } - PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), expected_raw_row_count)); + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), sealed_range.value())); } PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, merge_tree_writer_->PrepareCommit(wait_compaction)); @@ -285,7 +282,7 @@ Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compac } Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr& segment, - int64_t expected_raw_row_count) { + const OffsetRange& sealed_offsets) { PAIMON_ASSIGN_OR_RAISE(std::vector> readers, realtime_store_->CreateCommitReaders(segment)); ScopeGuard readers_guard([&readers]() { @@ -295,28 +292,25 @@ Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr> sorted_readers; - sorted_readers.reserve(readers.size()); - for (std::unique_ptr& reader : readers) { + for (const std::unique_ptr& reader : readers) { if (!reader) { return Status::Invalid("PK real-time store returned a null commit reader"); } - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr prepared_reader, - AdaptPreparedBatchReader(std::move(reader), prepared_schema_, std::nullopt, key_schema_, - write_schema_, memory_pool_, &raw_row_count)); + } + PAIMON_ASSIGN_OR_RAISE( + std::vector> prepared_readers, + AdaptPreparedCommitBatchReaders(std::move(readers), prepared_schema_, sealed_offsets, + key_schema_, write_schema_, key_comparator_, memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(prepared_readers.size()); + for (std::unique_ptr& prepared_reader : prepared_readers) { auto merge_function = std::make_unique(/*ignore_delete=*/false); sorted_readers.push_back(std::make_unique( std::move(prepared_reader), key_comparator_, std::make_shared(std::move(merge_function)))); } readers_guard.Release(); - PAIMON_RETURN_NOT_OK(merge_tree_writer_->WriteSortedReaders(std::move(sorted_readers))); - if (raw_row_count != expected_raw_row_count) { - return Status::Invalid("PK real-time store commit readers did not cover the sealed range"); - } - return Status::OK(); + return merge_tree_writer_->WriteSortedReaders(std::move(sorted_readers)); } Status RealtimePrimaryKeyWriter::Compact(bool) { diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index 9a5aa4c68..2eaf7ce24 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -79,7 +79,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { const std::shared_ptr& memory_pool); Status FlushSegment(const std::shared_ptr& segment, - int64_t expected_raw_row_count); + const OffsetRange& sealed_offsets); std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; diff --git a/src/paimon/core/realtime/realtime_reader.h b/src/paimon/core/realtime/realtime_reader.h index 6c25fd853..a041e0caa 100644 --- a/src/paimon/core/realtime/realtime_reader.h +++ b/src/paimon/core/realtime/realtime_reader.h @@ -44,10 +44,16 @@ class RealtimeReader final : public BatchReader { } Result NextBatch() override { + if (closed_) { + return MakeEofBatch(); + } return reader_->NextBatch(); } Result NextBatchWithBitmap() override { + if (closed_) { + return MakeEofBatchWithBitmap(); + } return reader_->NextBatchWithBitmap(); } @@ -56,6 +62,10 @@ class RealtimeReader final : public BatchReader { } void Close() override { + if (closed_) { + return; + } + closed_ = true; reader_->Close(); read_view_.reset(); } @@ -68,6 +78,7 @@ class RealtimeReader final : public BatchReader { // before releasing the data it references. std::shared_ptr read_view_; std::unique_ptr reader_; + bool closed_ = false; }; } // namespace paimon diff --git a/src/paimon/core/realtime/realtime_reader_test.cpp b/src/paimon/core/realtime/realtime_reader_test.cpp index ec37cfed4..10f6ce5be 100644 --- a/src/paimon/core/realtime/realtime_reader_test.cpp +++ b/src/paimon/core/realtime/realtime_reader_test.cpp @@ -37,6 +37,8 @@ class TestingReadView : public RealtimeReadView { class TestingBatchReader : public BatchReader { public: + explicit TestingBatchReader(int32_t* close_count = nullptr) : close_count_(close_count) {} + Result NextBatch() override { return MakeEofBatch(); } @@ -45,7 +47,14 @@ class TestingBatchReader : public BatchReader { return nullptr; } - void Close() override {} + void Close() override { + if (close_count_) { + ++(*close_count_); + } + } + + private: + int32_t* close_count_; }; TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { @@ -57,5 +66,21 @@ TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { "inner reader is null"); } +TEST(RealtimeReaderTest, TestCloseIsIdempotentAndReturnsEof) { + int32_t close_count = 0; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + RealtimeReader::Create(std::make_shared(), + std::make_unique(&close_count))); + reader->Close(); + reader->Close(); + ASSERT_EQ(1, close_count); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch_with_bitmap)); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index dc69ebb55..64d722097 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -91,11 +91,12 @@ Result>> CreateMemoryReaders( if (!reader) { return Status::Invalid("PK real-time store returned a null query reader"); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_reader, - AdaptPreparedBatchReader(std::move(reader), prepared_schema, - OffsetRange(split->CommittedEndOffset(), - split->MemoryEndOffset()), - key_schema, value_schema, memory_pool)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr prepared_reader, + AdaptPreparedBatchReader( + std::move(reader), prepared_schema, + OffsetRange(split->CommittedEndOffset(), split->MemoryEndOffset()), key_schema, + value_schema, key_comparator, memory_pool)); auto merge = std::make_unique(false); result.push_back(std::make_unique( std::move(prepared_reader), key_comparator, diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index dcf10e90c..0bcd79f61 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -233,7 +233,7 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& return Status::Invalid("real-time union read does not support data evolution"); } if (!table_schema.PrimaryKeys().empty()) { - PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options)); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options, table_schema)); } if (context.IsStreamingMode()) { return Status::Invalid("real-time union read currently supports batch scans only"); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index e68ff670d..e27b36903 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -423,6 +423,208 @@ class SplitCommitReaderRealtimeStoreFactory final : public RealtimeStoreFactory ArrowRealtimeStoreFactory delegate_; }; +class DropLastBatchReader final : public BatchReader { + public: + explicit DropLastBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + if (!buffered_.has_value()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch first, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(first)) { + return MakeEofBatch(); + } + buffered_ = std::move(first); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch next, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(next)) { + buffered_.reset(); + return MakeEofBatch(); + } + ReadBatch result = std::move(buffered_.value()); + buffered_ = std::move(next); + return result; + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + buffered_.reset(); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::optional buffered_; +}; + +class SwapFirstTwoBatchReader final : public BatchReader { + public: + explicit SwapFirstTwoBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + if (!initialized_) { + initialized_ = true; + PAIMON_ASSIGN_OR_RAISE(ReadBatch first, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(first)) { + return MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch second, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(second)) { + return first; + } + first_ = std::move(first); + return second; + } + if (first_.has_value()) { + ReadBatch first = std::move(first_.value()); + first_.reset(); + return first; + } + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + first_.reset(); + delegate_->Close(); + } + + private: + bool initialized_ = false; + std::unique_ptr delegate_; + std::optional first_; +}; + +class SubstituteOffsetBatchReader final : public BatchReader { + public: + explicit SubstituteOffsetBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT || array->length() == 0) { + return Status::Invalid("offset substitution requires a non-empty struct batch"); + } + std::shared_ptr struct_array = + std::dynamic_pointer_cast(array); + std::shared_ptr offsets = + std::dynamic_pointer_cast(struct_array->field(2)); + if (!offsets) { + return Status::Invalid("offset substitution requires an int64 REALTIME_OFFSET"); + } + arrow::Int64Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(offsets->length())); + for (int64_t row = 0; row < offsets->length(); ++row) { + builder.UnsafeAppend(0); + } + std::shared_ptr substituted_offsets; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&substituted_offsets)); + std::shared_ptr substituted_data = struct_array->data()->Copy(); + substituted_data->child_data[2] = substituted_offsets->data(); + std::shared_ptr substituted = arrow::MakeArray(std::move(substituted_data)); + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*substituted, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; +}; + +enum class CommitReaderMalformation { DROP_LAST, UNSORTED, SUBSTITUTE_OFFSET }; + +class MalformedCoverageRealtimeStore final : public RealtimeStore { + public: + MalformedCoverageRealtimeStore(const std::shared_ptr& delegate, + CommitReaderMalformation malformation) + : delegate_(delegate), malformation_(malformation) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + switch (malformation_) { + case CommitReaderMalformation::DROP_LAST: + reader = std::make_unique(std::move(reader)); + break; + case CommitReaderMalformation::UNSORTED: + reader = std::make_unique(std::move(reader)); + break; + case CommitReaderMalformation::SUBSTITUTE_OFFSET: + reader = std::make_unique(std::move(reader)); + break; + } + } + return readers; + } + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + CommitReaderMalformation malformation_; +}; + +class MalformedCoverageRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit MalformedCoverageRealtimeStoreFactory( + CommitReaderMalformation malformation = CommitReaderMalformation::DROP_LAST) + : malformation_(malformation) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate, malformation_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + CommitReaderMalformation malformation_; +}; + } // namespace namespace { @@ -1312,6 +1514,50 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { ASSERT_TRUE(query_view->expired()); } +TEST_F(RealtimeWriteInteTest, TestPkDeleteInsertAndPinnedReadsAcrossRefresh) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr delete_batch, + MakeBatch({Row{1, "deleted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(delete_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr insert_batch, + MakeBatch({Row{1, "inserted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(insert_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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 pinned_reader, + table_read->CreateReader(reader_plan->Splits())); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector plan_rows, ReadRows(pinned_plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "inserted", "p0"}}), plan_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_rows, + ReadResultCollector::CollectResult(pinned_reader.get())); + ASSERT_EQ(1, reader_rows->length()); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { options_[Options::READ_BATCH_SIZE] = "2"; CreatePkTable(); @@ -1903,6 +2149,56 @@ TEST_F(RealtimeWriteInteTest, TestPkPluginContract) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkRejectsMalformedCoverage) { + CreatePkTable(); + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "commit readers did not cover the sealed range"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsEqualCardinalityOffsetSubstitution) { + CreatePkTable(); + auto factory = std::make_shared( + CommitReaderMalformation::SUBSTITUTE_OFFSET); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "duplicate REALTIME_OFFSET"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsUnsortedPluginRows) { + CreatePkTable(); + auto factory = + std::make_shared(CommitReaderMalformation::UNSORTED); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "not globally sorted by primary key and sequence number"); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkQueryReaderClose) { CreatePkTable(); auto state = std::make_shared(); @@ -1973,10 +2269,10 @@ TEST_F(RealtimeWriteInteTest, TestPkQueryReaderCloseFailure) { return table_read->CreateReader(plan->Splits()); }; - for (int32_t null_index = 0; null_index <= 2; ++null_index) { + for (int32_t null_index = 0; null_index <= 1; ++null_index) { state->query_null_index = null_index; ASSERT_NOK_WITH_MSG(create_reader(), "PK real-time store returned a null query reader"); - ASSERT_EQ(2 * (null_index + 1), state->query_close_count->load(std::memory_order_acquire)); + ASSERT_EQ(null_index + 1, state->query_close_count->load(std::memory_order_acquire)); } ASSERT_OK(writer->Close()); } From a9aaa4e5ea98893c20cb15591a1675cf690cfc1c Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:27:48 +0800 Subject: [PATCH 21/24] refactor(realtime): simplify reader lifecycle cleanup --- .../merged_key_value_record_reader_test.cpp | 9 +------ .../key_value_file_store_write_test.cpp | 10 ++++--- .../realtime/arrow_realtime_store_test.cpp | 8 +++++- .../realtime/prepared_key_value_reader.cpp | 4 --- .../realtime/primary_key_realtime_store.cpp | 8 ------ .../primary_key_realtime_store_test.cpp | 3 --- .../realtime/realtime_append_only_writer.cpp | 2 +- .../core/realtime/realtime_context_impl.cpp | 26 +++++++++++++++---- .../core/realtime/realtime_context_impl.h | 3 +++ src/paimon/core/realtime/realtime_reader.h | 11 -------- .../core/realtime/realtime_reader_test.cpp | 15 +++++------ test/inte/realtime_write_inte_test.cpp | 2 +- 12 files changed, 48 insertions(+), 53 deletions(-) diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 775f271e3..a83c9bbd6 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -81,16 +81,11 @@ class TrackingBatchReader : public BatchReader { } void Close() override { - if (closed_) { - return; - } - closed_ = true; ++(*close_count_); delegate_->Close(); } private: - bool closed_ = false; std::unique_ptr delegate_; int32_t* close_count_; }; @@ -421,7 +416,7 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(1), 23); } -TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderClose) { +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderLifecycle) { std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), DataField(1, arrow::field("v0", arrow::int32()))}; std::shared_ptr value_schema = @@ -445,7 +440,6 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderClose) { AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), key_schema, value_schema, pool_)); reader->Close(); - reader->Close(); } ASSERT_EQ(explicit_close_count, 1); @@ -486,7 +480,6 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderClose) { key_schema, value_schema, pool_)); ASSERT_NOK_WITH_MSG(reader->NextBatch(), "prepared reader failure"); ASSERT_EQ(read_failure_close_count, 1); - reader->Close(); } ASSERT_EQ(read_failure_close_count, 1); } diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index 733c19d6e..a2344e803 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -411,6 +411,7 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { const std::map options = { {Options::BUCKET, "1"}, {Options::WRITE_BUFFER_SIZE, "1"}, + {Options::REALTIME_ENABLED, "true"}, }; const std::shared_ptr schema = arrow::schema({ arrow::field("id", arrow::int64(), false), @@ -465,7 +466,8 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { } TEST_F(KeyValueFileStoreWriteTest, TestRealtimeOffsetCollision) { - const std::map options = {{Options::BUCKET, "1"}}; + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; const std::shared_ptr schema = arrow::schema({ arrow::field("id", arrow::int64(), false), arrow::field("_REALTIME_OFFSET", arrow::int64()), @@ -488,7 +490,8 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeOffsetCollision) { } TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { - const std::map options = {{Options::BUCKET, "1"}}; + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; const std::shared_ptr schema = arrow::schema({ arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), @@ -546,7 +549,8 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { const int64_t max = std::numeric_limits::max(); - const std::map options = {{Options::BUCKET, "1"}}; + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; const std::shared_ptr schema = arrow::schema({ arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index 9aae99332..f186a8161 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -232,8 +232,14 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { ArrowRealtimeStoreFactory factory; std::unique_ptr write_schema = MakeReadSchema(schema_); + RealtimeStoreCreateRequest request{std::move(write_schema), + /*options=*/{}, + pool_, + /*partition=*/{}, + /*bucket=*/0, + AppendRealtimeStoreCreateConfig{StatisticsMode::FULL}}; ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_store, - factory.Create(std::move(write_schema), StatisticsMode::FULL, {}, pool_)); + factory.Create(std::move(request))); std::shared_ptr store = std::dynamic_pointer_cast(realtime_store); ASSERT_NE(nullptr, store); diff --git a/src/paimon/core/realtime/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp index 6b3afcd19..5b0375ad1 100644 --- a/src/paimon/core/realtime/prepared_key_value_reader.cpp +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -490,10 +490,6 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { private: Result> NextBatchImpl() { - if (closed_) { - return std::unique_ptr(); - } - while (true) { ResetBatchState(); PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index f43f60472..2f04aae79 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -177,9 +177,6 @@ class RawBatchReader final : public BatchReader { } Result NextBatch() override { - if (closed_) { - return MakeEofBatch(); - } std::optional selected; for (size_t i = 0; i < batches_.size(); ++i) { if (positions_[i] >= batches_[i].data->length()) { @@ -214,10 +211,6 @@ class RawBatchReader final : public BatchReader { return metrics_; } void Close() override { - if (closed_) { - return; - } - closed_ = true; batches_.clear(); positions_.clear(); key_contexts_.clear(); @@ -238,7 +231,6 @@ class RawBatchReader final : public BatchReader { return left_sequences->Value(positions_[left]) < right_sequences->Value(positions_[right]); } - bool closed_ = false; std::vector batches_; std::vector positions_; std::vector key_field_indexes_; diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index cafe3682e..116c6e389 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -218,9 +218,6 @@ TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { "4 type: string\n [\n \"before\",\n \"after\",\n \"three\"\n ]", actual); readers[0]->Close(); - readers[0]->Close(); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, readers[0]->NextBatch()); - ASSERT_TRUE(BatchReader::IsEofBatch(eof)); } TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderExportsZeroOffsets) { diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 21d6cfb74..ea5feecce 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -47,7 +47,7 @@ Result> RealtimeAppendOnlyWriter::Crea std::unique_ptr<::ArrowSchema> write_schema, const std::shared_ptr& realtime_context, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, + const std::shared_ptr& input_schema, StatisticsMode statistics_mode, const std::map& options, const std::shared_ptr& memory_pool) { if (!realtime_context) { diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 215e066ee..ba4c8b7a6 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -269,12 +269,28 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, if (partition_bucket.bucket < 0 || committed_end_offset < 0) { return Status::Invalid("invalid partition-bucket committed offset"); } + } + // Only stores created by this context can contain state which cannot be restored in + // place. Offsets for other partition-buckets are reference state for lazy store creation + // and may be removed or rolled back without rebuilding the context. + std::lock_guard registry_lock(mutex_); + for (const auto& store_entry : stores_) { + const RealtimePartitionBucket& partition_bucket = store_entry.first; auto previous_iter = committed_offsets_.find(partition_bucket); - if (previous_iter != committed_offsets_.end()) { - if (committed_end_offset < previous_iter->second) { - return Status::Invalid( - "real-time partition-bucket committed offset cannot move backwards"); - } + if (previous_iter == committed_offsets_.end()) { + continue; + } + + auto current_iter = committed_offsets.find(partition_bucket); + if (current_iter == committed_offsets.end()) { + return Status::Invalid( + "real-time committed progress removed an active partition-bucket; recreate " + "RealtimeContext"); + } + if (current_iter->second < previous_iter->second) { + return Status::Invalid( + "real-time committed offset moved backwards for an active partition-bucket; " + "recreate RealtimeContext"); } } committed_offsets_ = committed_offsets; diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 9fa145e99..f5118c18f 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -88,6 +88,9 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { Status ReleaseReadView(const std::string& opaque_ticket); + // Returns an error requiring a new context if a newer snapshot removes or moves committed + // progress backwards for a store created by this context. Progress for inactive stores is + // only reference state and can be replaced in place. Status AdvanceCommittedProgress(int64_t snapshot_id, const RealtimeOffsetMap& committed_offsets); diff --git a/src/paimon/core/realtime/realtime_reader.h b/src/paimon/core/realtime/realtime_reader.h index a041e0caa..6c25fd853 100644 --- a/src/paimon/core/realtime/realtime_reader.h +++ b/src/paimon/core/realtime/realtime_reader.h @@ -44,16 +44,10 @@ class RealtimeReader final : public BatchReader { } Result NextBatch() override { - if (closed_) { - return MakeEofBatch(); - } return reader_->NextBatch(); } Result NextBatchWithBitmap() override { - if (closed_) { - return MakeEofBatchWithBitmap(); - } return reader_->NextBatchWithBitmap(); } @@ -62,10 +56,6 @@ class RealtimeReader final : public BatchReader { } void Close() override { - if (closed_) { - return; - } - closed_ = true; reader_->Close(); read_view_.reset(); } @@ -78,7 +68,6 @@ class RealtimeReader final : public BatchReader { // before releasing the data it references. std::shared_ptr read_view_; std::unique_ptr reader_; - bool closed_ = false; }; } // namespace paimon diff --git a/src/paimon/core/realtime/realtime_reader_test.cpp b/src/paimon/core/realtime/realtime_reader_test.cpp index 10f6ce5be..ded060989 100644 --- a/src/paimon/core/realtime/realtime_reader_test.cpp +++ b/src/paimon/core/realtime/realtime_reader_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include "paimon/arrow/abi.h" #include "paimon/testing/utils/testharness.h" @@ -66,20 +67,18 @@ TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { "inner reader is null"); } -TEST(RealtimeReaderTest, TestCloseIsIdempotentAndReturnsEof) { +TEST(RealtimeReaderTest, TestCloseReleasesResources) { int32_t close_count = 0; + std::shared_ptr read_view = std::make_shared(); + std::weak_ptr weak_read_view = read_view; ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, - RealtimeReader::Create(std::make_shared(), + RealtimeReader::Create(std::move(read_view), std::make_unique(&close_count))); - reader->Close(); + ASSERT_FALSE(weak_read_view.expired()); reader->Close(); ASSERT_EQ(1, close_count); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); - ASSERT_TRUE(BatchReader::IsEofBatch(batch)); - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, - reader->NextBatchWithBitmap()); - ASSERT_TRUE(BatchReader::IsEofBatch(batch_with_bitmap)); + ASSERT_TRUE(weak_read_view.expired()); } } // namespace diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index e27b36903..0e6f83b70 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1939,7 +1939,7 @@ TEST_F(RealtimeWriteInteTest, TestPkRecovery) { seed_commit_builder.SetOptions(options_).Finish()); ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit, FileStoreCommit::Create(std::move(seed_commit_context))); - ASSERT_OK(seed_commit->Commit(seed_messages)); + ASSERT_OK(seed_commit->Commit(seed_messages, /*commit_identifier=*/0)); ASSERT_OK(seed_writer->Close()); const std::vector mutations = { {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; From 8ab981752c64242cee80cfab66d99c17dbd2febd Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:01:26 +0800 Subject: [PATCH 22/24] fix(realtime): harden primary-key prepared batches --- include/paimon/realtime/realtime_store.h | 2 + include/paimon/utils/special_field_ids.h | 2 + .../io/merged_key_value_record_reader.cpp | 10 +- .../core/io/merged_key_value_record_reader.h | 1 + .../merged_key_value_record_reader_test.cpp | 93 ++++++++- src/paimon/core/mergetree/merge_tree_writer.h | 3 + .../realtime/prepared_key_value_reader.cpp | 144 ++++++++------ .../core/realtime/prepared_key_value_reader.h | 2 + .../realtime/primary_key_realtime_store.cpp | 135 ++++++++++--- .../primary_key_realtime_store_test.cpp | 128 ++++++++++++- src/paimon/core/realtime/realtime_fields.h | 6 +- .../realtime/realtime_primary_key_writer.cpp | 14 -- .../table/source/append_only_table_read.cpp | 37 +++- .../table/source/key_value_table_read.cpp | 12 +- .../core/table/source/realtime_table_scan.cpp | 20 +- .../core/table/source/realtime_table_scan.h | 3 +- src/paimon/core/table/source/table_scan.cpp | 2 +- test/inte/realtime_write_inte_test.cpp | 180 ++++++++++++++++-- 18 files changed, 652 insertions(+), 142 deletions(-) diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index 792bb1c56..90c6ce0a8 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -109,6 +109,8 @@ class PAIMON_EXPORT RealtimeReadView { struct PAIMON_EXPORT RealtimeQueryContext { /// Append mode receives the requested output fields before the mandatory leading /// `_VALUE_KIND` field is added. Primary-key mode receives the complete prepared schema. + /// This schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must + /// import or copy it synchronously. ::ArrowSchema* read_schema; /// Predicate using field indexes from `read_schema`. std::shared_ptr predicate; diff --git a/include/paimon/utils/special_field_ids.h b/include/paimon/utils/special_field_ids.h index 829f29889..5219d72db 100644 --- a/include/paimon/utils/special_field_ids.h +++ b/include/paimon/utils/special_field_ids.h @@ -42,6 +42,8 @@ class SpecialFieldIds { /// Special field ID reserved for index score. Value: CPP_FIELD_ID_END - 1 inline static constexpr int32_t INDEX_SCORE = CPP_FIELD_ID_END - 1; + /// Special field ID reserved for realtime offset. Value: CPP_FIELD_ID_END - 2 + inline static constexpr int32_t REALTIME_OFFSET = CPP_FIELD_ID_END - 2; /// Lowest field ID reserved for system fields; IDs at or above it are excluded from the /// highest field ID of a schema. Value: INT32_MAX / 2 diff --git a/src/paimon/core/io/merged_key_value_record_reader.cpp b/src/paimon/core/io/merged_key_value_record_reader.cpp index 70f2bcfb9..8c3952874 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader.cpp @@ -117,13 +117,21 @@ Status MergedKeyValueRecordReader::Iterator::LoadNextRawKeyValue() const { } Result> MergedKeyValueRecordReader::NextBatch() { + if (initialization_error_.has_value()) { + return initialization_error_.value(); + } if (visited_) { return std::unique_ptr(); } visited_ = true; auto iterator = std::make_unique(this); - PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); + Result has_next_result = iterator->HasNext(); + if (!has_next_result.ok()) { + initialization_error_ = has_next_result.status(); + return initialization_error_.value(); + } + bool has_next = std::move(has_next_result).value(); if (!has_next) { return std::unique_ptr(); } diff --git a/src/paimon/core/io/merged_key_value_record_reader.h b/src/paimon/core/io/merged_key_value_record_reader.h index a1b7aa5e4..227a1593a 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.h +++ b/src/paimon/core/io/merged_key_value_record_reader.h @@ -67,6 +67,7 @@ class MergedKeyValueRecordReader : public KeyValueRecordReader { private: bool visited_ = false; + std::optional initialization_error_; std::unique_ptr reader_; std::shared_ptr key_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index a83c9bbd6..a0d65205c 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -19,7 +19,6 @@ #include "paimon/core/io/merged_key_value_record_reader.h" #include -#include #include #include #include @@ -45,6 +44,7 @@ #include "paimon/testing/utils/key_value_checker.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/utils/special_field_ids.h" namespace paimon::test { @@ -107,7 +107,7 @@ class MergedKeyValueRecordReaderTest : public testing::Test { TEST_F(MergedKeyValueRecordReaderTest, TestRealtimeOffsetField) { const DataField& field = RealtimeOffsetField(); - ASSERT_EQ(std::numeric_limits::max() - 10002, field.Id()); + ASSERT_EQ(SpecialFieldIds::REALTIME_OFFSET, field.Id()); ASSERT_EQ("_REALTIME_OFFSET", field.Name()); ASSERT_EQ(arrow::Type::INT64, field.Type()->id()); ASSERT_FALSE(field.Nullable()); @@ -296,6 +296,95 @@ TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { "exact"); } +TEST_F(MergedKeyValueRecordReaderTest, TestBadCommitBatch) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value = MakeField("value", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, std::nullopt, + arrow::schema({key}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "field count"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMissingCompositeKey) { + std::shared_ptr key0 = MakeField("key0", arrow::int32(), 0); + std::shared_ptr key1 = MakeField("key1", arrow::int32(), 1); + std::shared_ptr value = MakeField("value", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key0, key1, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key0, key1, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key0, value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key0, key1}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "cannot find field id 1"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestQueryAddRename) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr old_value = MakeField("old_value", arrow::int32(), 1); + std::shared_ptr renamed_value = MakeField("renamed_value", arrow::int32(), 1); + std::shared_ptr added = MakeField("added", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key, renamed_value, added}); + std::shared_ptr prepared_schema = + MakePreparedSchema({key, renamed_value, added}); + std::shared_ptr actual_schema = MakePreparedSchema({key, old_value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key}), value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr iterator, + reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(KeyValue key_value, iterator->Next()); + ASSERT_EQ(20, key_value.value->GetInt(1)); + ASSERT_TRUE(key_value.value->IsNullAt(2)); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMergedReaderErrorRetry) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + std::shared_ptr prepared_array = + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + auto failing_reader = std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("stable prepared error")); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(failing_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({DataField(0, key)}, true)); + MergedKeyValueRecordReader merged_reader(std::move(reader), key_comparator, + merge_function_wrapper_); + + Result> first = merged_reader.NextBatch(); + Result> retry = merged_reader.NextBatch(); + ASSERT_NOK(first); + ASSERT_NOK(retry); + ASSERT_EQ(first.status().ToString(), retry.status().ToString()); +} + TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderSafeDecode) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index 542affd81..cea07f3e4 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -70,6 +70,9 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; + /// Consumes readers whose complete streams are individually sorted by primary key and + /// sequence number. Readers are closed on success or failure; an error may leave generated + /// file state unpublished, so the caller must discard this writer and replay its input. Status WriteSortedReaders(std::vector>&& readers); Status Compact(bool full_compaction) override; diff --git a/src/paimon/core/realtime/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp index 5b0375ad1..864456818 100644 --- a/src/paimon/core/realtime/prepared_key_value_reader.cpp +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -62,8 +62,18 @@ constexpr int32_t kSequenceNumberIndex = 1; constexpr int32_t kRealtimeOffsetIndex = 2; constexpr int32_t kPreparedValueStartIndex = 3; +template +void CloseReaders(const std::vector>& readers) { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + Result> AlignArrayByPaimonIds( - const std::shared_ptr& array, const std::shared_ptr& read_type); + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool); class RealtimeOffsetCoverage { public: @@ -237,22 +247,9 @@ Status ValidateExactCommitSchema(const std::shared_ptr& prepared_ return Status::OK(); } -Status ValidatePreparedSchema(const std::shared_ptr& prepared_schema) { - if (!prepared_schema || prepared_schema->num_fields() < kPreparedValueStartIndex) { - return Status::Invalid("prepared schema must contain realtime transport fields"); - } - PAIMON_RETURN_NOT_OK( - CheckPreparedField(prepared_schema, kValueKindIndex, SpecialFields::ValueKind())); - PAIMON_RETURN_NOT_OK( - CheckPreparedField(prepared_schema, kSequenceNumberIndex, SpecialFields::SequenceNumber())); - PAIMON_RETURN_NOT_OK( - CheckPreparedField(prepared_schema, kRealtimeOffsetIndex, RealtimeOffsetField())); - return Status::OK(); -} - Result> AlignStructArrayByPaimonIds( const std::shared_ptr& array, - const std::shared_ptr& read_type) { + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { const std::shared_ptr data_type = checked_pointer_cast(array->type()); std::unordered_map data_field_id_to_idx; @@ -273,12 +270,16 @@ Result> AlignStructArrayByPaimonIds( NestedProjectionUtils::GetPaimonFieldId(read_field)); auto data_iter = data_field_id_to_idx.find(read_field_id); if (data_iter == data_field_id_to_idx.end()) { - return Status::Invalid( - fmt::format("cannot find field id {} in prepared value struct", read_field_id)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr null_child, + arrow::MakeArrayOfNull(read_field->type(), array->offset() + array->length(), + arrow_pool)); + aligned_arrays.push_back(std::move(null_child)); + continue; } std::shared_ptr child = arrow::MakeArray(array->data()->child_data[data_iter->second]); - PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type())); + PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type(), arrow_pool)); aligned_arrays.push_back(std::move(child)); } @@ -294,9 +295,10 @@ Result> AlignStructArrayByPaimonIds( Result> AlignListArrayByPaimonIds( const std::shared_ptr& array, - const std::shared_ptr& read_type) { + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { std::shared_ptr values = array->values(); - PAIMON_ASSIGN_OR_RAISE(values, AlignArrayByPaimonIds(values, read_type->value_type())); + PAIMON_ASSIGN_OR_RAISE(values, + AlignArrayByPaimonIds(values, read_type->value_type(), arrow_pool)); std::shared_ptr new_data = array->data()->Copy(); new_data->type = read_type; new_data->child_data = {values->data()}; @@ -304,12 +306,12 @@ Result> AlignListArrayByPaimonIds( } Result> AlignMapArrayByPaimonIds( - const std::shared_ptr& array, - const std::shared_ptr& read_type) { + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { std::shared_ptr keys = array->keys(); - PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type())); + PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type(), arrow_pool)); std::shared_ptr items = array->items(); - PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type())); + PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type(), arrow_pool)); const std::shared_ptr& entries_data = array->data()->child_data[0]; std::shared_ptr new_entries = entries_data->Copy(); @@ -323,7 +325,8 @@ Result> AlignMapArrayByPaimonIds( } Result> AlignArrayByPaimonIds( - const std::shared_ptr& array, const std::shared_ptr& read_type) { + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { if (array->type()->id() != read_type->id()) { return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", array->type()->ToString(), read_type->ToString())); @@ -331,13 +334,16 @@ Result> AlignArrayByPaimonIds( switch (read_type->id()) { case arrow::Type::STRUCT: return AlignStructArrayByPaimonIds(checked_pointer_cast(array), - checked_pointer_cast(read_type)); + checked_pointer_cast(read_type), + arrow_pool); case arrow::Type::LIST: return AlignListArrayByPaimonIds(checked_pointer_cast(array), - checked_pointer_cast(read_type)); + checked_pointer_cast(read_type), + arrow_pool); case arrow::Type::MAP: return AlignMapArrayByPaimonIds(checked_pointer_cast(array), - checked_pointer_cast(read_type)); + checked_pointer_cast(read_type), + arrow_pool); default: if (!array->type()->Equals(*read_type)) { return Status::Invalid( @@ -351,7 +357,7 @@ Result> AlignArrayByPaimonIds( Result ProjectFieldsByPaimonIds( const std::shared_ptr& data_batch, const std::shared_ptr& prepared_schema, - const std::shared_ptr& query_schema) { + const std::shared_ptr& query_schema, arrow::MemoryPool* arrow_pool) { std::unordered_map prepared_field_id_to_idx; prepared_field_id_to_idx.reserve(prepared_schema->num_fields()); for (int32_t i = kPreparedValueStartIndex; i < prepared_schema->num_fields(); ++i) { @@ -375,7 +381,7 @@ Result ProjectFieldsByPaimonIds( } std::shared_ptr field_array = data_batch->field(prepared_iter->second); PAIMON_ASSIGN_OR_RAISE(field_array, - AlignArrayByPaimonIds(field_array, query_field->type())); + AlignArrayByPaimonIds(field_array, query_field->type(), arrow_pool)); result.push_back(std::move(field_array)); } return result; @@ -468,8 +474,12 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { }; Result> NextBatch() override { + if (first_error_.has_value()) { + return first_error_.value(); + } Result> result = NextBatchImpl(); if (!result.ok()) { + first_error_ = result.status(); Close(); } return result; @@ -508,6 +518,23 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { } std::shared_ptr data_batch = checked_pointer_cast(arrow_array); + Status transport_status = + ValidatePreparedTransportSchema(arrow::schema(data_batch->type()->fields())); + if (!transport_status.ok()) { + return Status::Invalid( + "prepared batch field does not match prepared transport " + "schema: ", + transport_status.ToString()); + } + if (visible_offsets_.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema( + arrow::schema(data_batch->type()->fields()), key_schema_)); + PAIMON_ASSIGN_OR_RAISE( + arrow_array, + AlignArrayByPaimonIds(data_batch, arrow::struct_(prepared_schema_->fields()), + arrow_pool_.get())); + data_batch = checked_pointer_cast(arrow_array); + } PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch)); PAIMON_RETURN_NOT_OK(ValidateOrdering(data_batch)); @@ -528,12 +555,12 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { data_batch->field(kValueKindIndex)); sequence_number_array_ = checked_pointer_cast>( data_batch->field(kSequenceNumberIndex)); - PAIMON_ASSIGN_OR_RAISE( - arrow::ArrayVector key_fields, - ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_)); - PAIMON_ASSIGN_OR_RAISE( - arrow::ArrayVector value_fields, - ProjectFieldsByPaimonIds(data_batch, prepared_schema_, value_schema_)); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + key_schema_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector value_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + value_schema_, arrow_pool_.get())); key_ctx_ = std::make_shared(key_fields, pool_); value_ctx_ = std::make_shared(value_fields, pool_); ArrowUtils::TraverseArray(data_batch); @@ -578,8 +605,9 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { if (data_batch->length() == 0) { return Status::OK(); } - PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields, - ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_)); + PAIMON_ASSIGN_OR_RAISE( + arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_, arrow_pool_.get())); std::shared_ptr key_context = std::make_shared(key_fields, pool_); std::shared_ptr sequences = @@ -613,6 +641,7 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { private: bool closed_ = false; + std::optional first_error_; std::unique_ptr reader_; std::shared_ptr prepared_schema_; std::optional visible_offsets_; @@ -634,6 +663,19 @@ class PreparedKeyValueReader final : public KeyValueRecordReader { } // namespace +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema) { + if (!prepared_schema || prepared_schema->num_fields() < kPreparedValueStartIndex) { + return Status::Invalid("prepared schema must contain realtime transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kSequenceNumberIndex, SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kRealtimeOffsetIndex, RealtimeOffsetField())); + return Status::OK(); +} + namespace { Result> AdaptPreparedBatchReaderImpl( @@ -649,7 +691,7 @@ Result> AdaptPreparedBatchReaderImpl( return Status::Invalid("prepared batch reader cannot be null"); } ScopeGuard close_guard([&owned_reader]() -> void { owned_reader->Close(); }); - PAIMON_RETURN_NOT_OK(ValidatePreparedSchema(prepared_schema)); + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); if (!key_schema) { return Status::Invalid("prepared key schema cannot be null"); } @@ -695,26 +737,23 @@ Result>> AdaptPreparedCommitBa const std::shared_ptr& value_schema, const std::shared_ptr& key_comparator, const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard readers_guard([&readers, &adapted_readers]() { + CloseReaders(readers); + CloseReaders(adapted_readers); + }); if (!memory_pool) { return Status::Invalid("prepared reader memory pool cannot be null"); } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } std::shared_ptr arrow_pool = GetArrowPool(memory_pool); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr offset_coverage, RealtimeOffsetCoverage::Create(sealed_offsets, readers.size(), arrow_pool)); - ScopeGuard readers_guard([&readers]() { - for (const std::unique_ptr& reader : readers) { - if (reader) { - reader->Close(); - } - } - }); - std::vector> adapted_readers; - ScopeGuard adapted_readers_guard([&adapted_readers]() { - for (const std::unique_ptr& reader : adapted_readers) { - reader->Close(); - } - }); adapted_readers.reserve(readers.size()); for (std::unique_ptr& reader : readers) { PAIMON_ASSIGN_OR_RAISE(std::unique_ptr adapted_reader, @@ -724,7 +763,6 @@ Result>> AdaptPreparedCommitBa adapted_readers.push_back(std::move(adapted_reader)); } readers_guard.Release(); - adapted_readers_guard.Release(); return adapted_readers; } diff --git a/src/paimon/core/realtime/prepared_key_value_reader.h b/src/paimon/core/realtime/prepared_key_value_reader.h index 064a62958..22a837a76 100644 --- a/src/paimon/core/realtime/prepared_key_value_reader.h +++ b/src/paimon/core/realtime/prepared_key_value_reader.h @@ -33,6 +33,8 @@ class BatchReader; class FieldsComparator; class MemoryPool; +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema); + Result> AdaptPreparedBatchReader( std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, const std::optional& visible_offsets, diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp index 2f04aae79..e4c480377 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include #include @@ -37,6 +39,7 @@ #include "paimon/common/utils/fields_comparator.h" #include "paimon/core/core_options.h" #include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" #include "paimon/core/schema/table_schema.h" #include "paimon/macros.h" @@ -163,9 +166,12 @@ class RawBatchReader final : public BatchReader { key_comparator_(key_comparator), memory_pool_(memory_pool), arrow_pool_(GetArrowPool(memory_pool)), + heap_(SourceGreater{this}), metrics_(std::make_shared()) { key_contexts_.reserve(batches_.size()); - for (const StoredBatch& batch : batches_) { + sequence_arrays_.reserve(batches_.size()); + for (size_t i = 0; i < batches_.size(); ++i) { + const StoredBatch& batch = batches_[i]; arrow::ArrayVector key_arrays; key_arrays.reserve(key_field_indexes_.size()); for (int32_t field_index : key_field_indexes_) { @@ -173,34 +179,89 @@ class RawBatchReader final : public BatchReader { } key_contexts_.push_back( std::make_shared(key_arrays, memory_pool_)); + sequence_arrays_.push_back( + checked_pointer_cast(batch.data->field(1))); + if (batch.data->length() > 0) { + heap_.push(i); + } } } Result NextBatch() override { - std::optional selected; - for (size_t i = 0; i < batches_.size(); ++i) { - if (positions_[i] >= batches_[i].data->length()) { - continue; + if (heap_.empty()) { + return MakeEofBatch(); + } + + struct SelectedRow { + size_t selected_source; + int64_t source_ordinal; + }; + struct SelectedSource { + size_t source; + std::vector rows; + int64_t base = -1; + }; + std::vector selected_rows; + selected_rows.reserve(kOutputBatchSize); + std::vector selected_sources; + std::unordered_map selected_source_indexes; + while (!heap_.empty() && selected_rows.size() < kOutputBatchSize) { + const size_t source = heap_.top(); + heap_.pop(); + auto [source_it, inserted] = + selected_source_indexes.emplace(source, selected_sources.size()); + if (inserted) { + selected_sources.push_back(SelectedSource{source, {}}); } - if (!selected.has_value() || Less(i, selected.value())) { - selected = i; + SelectedSource& selected_source = selected_sources[source_it->second]; + selected_rows.push_back( + SelectedRow{source_it->second, static_cast(selected_source.rows.size())}); + selected_source.rows.push_back(positions_[source]++); + if (positions_[source] < batches_[source].data->length()) { + heap_.push(source); } } - if (!selected.has_value()) { - return MakeEofBatch(); - } - const size_t batch_index = selected.value(); - arrow::Int64Builder index_builder(arrow_pool_.get()); - PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Append(positions_[batch_index])); - std::shared_ptr index; - PAIMON_RETURN_NOT_OK_FROM_ARROW(index_builder.Finish(&index)); + arrow::compute::ExecContext context(arrow_pool_.get()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - arrow::Datum taken, - arrow::compute::Take(arrow::Datum(batches_[batch_index].data), arrow::Datum(index), - arrow::compute::TakeOptions::NoBoundsCheck(), &context)); - std::shared_ptr batch = taken.make_array(); - ++positions_[batch_index]; + arrow::ArrayVector grouped_batches; + int64_t grouped_row_count = 0; + for (SelectedSource& selected_source : selected_sources) { + arrow::Int64Builder source_index_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + source_index_builder.AppendValues(selected_source.rows)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr source_indices, + source_index_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum source_batch, + arrow::compute::Take(arrow::Datum(batches_[selected_source.source].data), + arrow::Datum(source_indices), + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + selected_source.base = grouped_row_count; + grouped_row_count += static_cast(selected_source.rows.size()); + grouped_batches.push_back(source_batch.make_array()); + } + + std::shared_ptr batch; + if (grouped_batches.size() == 1) { + batch = std::move(grouped_batches[0]); + } else { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr grouped, + arrow::Concatenate(grouped_batches, arrow_pool_.get())); + arrow::Int64Builder order_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(order_builder.Reserve(selected_rows.size())); + for (const SelectedRow& selected : selected_rows) { + order_builder.UnsafeAppend(selected_sources[selected.selected_source].base + + selected.source_ordinal); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr order, + order_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum reordered, + arrow::compute::Take(arrow::Datum(grouped), arrow::Datum(order), + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + batch = reordered.make_array(); + } auto array = std::make_unique(); auto schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, array.get(), schema.get())); @@ -211,12 +272,18 @@ class RawBatchReader final : public BatchReader { return metrics_; } void Close() override { + while (!heap_.empty()) { + heap_.pop(); + } batches_.clear(); positions_.clear(); key_contexts_.clear(); + sequence_arrays_.clear(); } private: + static constexpr size_t kOutputBatchSize = 1024; + bool Less(size_t left, size_t right) const { ColumnarRowRef left_key(key_contexts_[left], positions_[left]); ColumnarRowRef right_key(key_contexts_[right], positions_[right]); @@ -224,13 +291,22 @@ class RawBatchReader final : public BatchReader { if (key_comparison != 0) { return key_comparison < 0; } - const std::shared_ptr left_sequences = - checked_pointer_cast(batches_[left].data->field(1)); - const std::shared_ptr right_sequences = - checked_pointer_cast(batches_[right].data->field(1)); - return left_sequences->Value(positions_[left]) < right_sequences->Value(positions_[right]); + const int64_t left_sequence = sequence_arrays_[left]->Value(positions_[left]); + const int64_t right_sequence = sequence_arrays_[right]->Value(positions_[right]); + if (left_sequence != right_sequence) { + return left_sequence < right_sequence; + } + return left < right; } + struct SourceGreater { + RawBatchReader* reader; + + bool operator()(size_t left, size_t right) const { + return reader->Less(right, left); + } + }; + std::vector batches_; std::vector positions_; std::vector key_field_indexes_; @@ -238,6 +314,8 @@ class RawBatchReader final : public BatchReader { std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; std::vector> key_contexts_; + std::vector> sequence_arrays_; + std::priority_queue, SourceGreater> heap_; std::shared_ptr metrics_; }; @@ -379,8 +457,9 @@ Result> PrimaryKeyRealtimeStore::Create const std::shared_ptr& prepared_schema, const std::vector& trimmed_primary_keys, const std::shared_ptr& memory_pool) { - if (!prepared_schema || trimmed_primary_keys.empty() || !memory_pool) { - return Status::Invalid("PK prepared schema or memory pool is null"); + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); + if (trimmed_primary_keys.empty() || !memory_pool) { + return Status::Invalid("PK primary keys are empty or memory pool is null"); } std::vector key_field_indexes; std::vector key_fields; diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index 116c6e389..dc2ce86b1 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -18,14 +18,18 @@ #include "paimon/core/realtime/primary_key_realtime_store.h" +#include #include +#include #include #include +#include #include #include "arrow/api.h" #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" +#include "fmt/format.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -196,6 +200,34 @@ TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); } +TEST(PrimaryKeyRealtimeStoreTest, TestBadTransportPrefix) { + const std::shared_ptr valid = PreparedSchema(); + std::vector invalid_fields; + + arrow::FieldVector wrong_type = valid->fields(); + wrong_type[0] = DataField::ConvertDataFieldToArrowField( + DataField(SpecialFields::ValueKind().Id(), + arrow::field("_VALUE_KIND", arrow::int32(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_type)); + + arrow::FieldVector nullable_sequence = valid->fields(); + nullable_sequence[1] = nullable_sequence[1]->WithNullable(true); + invalid_fields.push_back(std::move(nullable_sequence)); + + arrow::FieldVector wrong_offset_id = valid->fields(); + wrong_offset_id[2] = DataField::ConvertDataFieldToArrowField( + DataField(99, arrow::field("_REALTIME_OFFSET", arrow::int64(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_offset_id)); + + for (const arrow::FieldVector& fields : invalid_fields) { + ASSERT_NOK_WITH_MSG( + PrimaryKeyRealtimeStore::Create(arrow::schema(fields), {"id"}, GetDefaultPool()), + "prepared schema field"); + } +} + TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { ASSERT_OK_AND_ASSIGN( std::shared_ptr store, @@ -233,12 +265,100 @@ TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderExportsZeroOffsets) { ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateCommitReaders(segment.value())); ASSERT_EQ(1, readers.size()); - for (int32_t row = 0; row < 2; ++row) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + ASSERT_TRUE(arrow::ImportArray(batch.first.get(), batch.second.get()).ok()); + ASSERT_OK_AND_ASSIGN(batch, readers[0]->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestHeapMergeAcrossBatches) { + constexpr int64_t kSourceCount = 2057; + constexpr int64_t kKeyCount = 257; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + for (int64_t source = 0; source < kSourceCount; ++source) { + const int64_t id = (source * 149) % kKeyCount; + const std::string json = + fmt::format(R"([[0, {}, {}, {}, "v{}"]])", source, source, id, source); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(json), OffsetRange(source, source + 1)})); + } + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + + std::vector expected_sources(kSourceCount); + std::iota(expected_sources.begin(), expected_sources.end(), 0); + std::sort(expected_sources.begin(), expected_sources.end(), [=](int64_t left, int64_t right) { + const int64_t left_id = (left * 149) % kKeyCount; + const int64_t right_id = (right * 149) % kKeyCount; + return left_id != right_id ? left_id < right_id : left < right; + }); + + int64_t output_row = 0; + int64_t output_batches = 0; + while (true) { ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); - ASSERT_FALSE(BatchReader::IsEofBatch(batch)); - AssertOffsetsZero(batch.first.get()); - ASSERT_TRUE(arrow::ImportArray(batch.first.get(), batch.second.get()).ok()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + ASSERT_LE(batch.first->length, 1024); + ASSERT_GT(batch.first->length, 0); + ++output_batches; + arrow::Result> imported_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported_result.ok()) << imported_result.status().ToString(); + std::shared_ptr imported = std::move(imported_result).ValueOrDie(); + std::shared_ptr array = + std::dynamic_pointer_cast(imported); + ASSERT_NE(nullptr, array); + ASSERT_EQ(PreparedSchema()->ToString(), arrow::schema(array->type()->fields())->ToString()); + std::shared_ptr sequences = + std::dynamic_pointer_cast(array->field(1)); + std::shared_ptr ids = + std::dynamic_pointer_cast(array->field(3)); + std::shared_ptr values = + std::dynamic_pointer_cast(array->field(4)); + ASSERT_NE(nullptr, sequences); + ASSERT_NE(nullptr, ids); + ASSERT_NE(nullptr, values); + for (int64_t row = 0; row < array->length(); ++row, ++output_row) { + ASSERT_LT(output_row, kSourceCount); + const int64_t source = expected_sources[output_row]; + ASSERT_EQ(source, sequences->Value(row)); + ASSERT_EQ((source * 149) % kKeyCount, ids->Value(row)); + ASSERT_EQ(fmt::format("v{}", source), values->GetString(row)); + } } + ASSERT_EQ(kSourceCount, output_row); + ASSERT_EQ(3, output_batches); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCloseUnreadMultiSourceReader) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 10, 0, 1, "a"]])"), OffsetRange(0, 1)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 20, 1, 2, "b"]])"), OffsetRange(1, 2)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 30, 2, 3, "c"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + + readers[0]->Close(); } TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { diff --git a/src/paimon/core/realtime/realtime_fields.h b/src/paimon/core/realtime/realtime_fields.h index 6ed04b38a..270941238 100644 --- a/src/paimon/core/realtime/realtime_fields.h +++ b/src/paimon/core/realtime/realtime_fields.h @@ -19,17 +19,15 @@ #pragma once -#include -#include - #include "arrow/type.h" #include "paimon/common/types/data_field.h" +#include "paimon/utils/special_field_ids.h" namespace paimon { inline const DataField& RealtimeOffsetField() { static const DataField data_field = - DataField(std::numeric_limits::max() - 10002, + DataField(SpecialFieldIds::REALTIME_OFFSET, arrow::field("_REALTIME_OFFSET", arrow::int64(), /*nullable=*/false)); return data_field; } diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 82318eadb..185156eb7 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -32,7 +32,6 @@ #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/scope_guard.h" #include "paimon/core/io/merged_key_value_record_reader.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -285,18 +284,6 @@ Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr> readers, realtime_store_->CreateCommitReaders(segment)); - ScopeGuard readers_guard([&readers]() { - for (const std::unique_ptr& reader : readers) { - if (reader) { - reader->Close(); - } - } - }); - for (const std::unique_ptr& reader : readers) { - if (!reader) { - return Status::Invalid("PK real-time store returned a null commit reader"); - } - } PAIMON_ASSIGN_OR_RAISE( std::vector> prepared_readers, AdaptPreparedCommitBatchReaders(std::move(readers), prepared_schema_, sealed_offsets, @@ -309,7 +296,6 @@ Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr(std::move(merge_function)))); } - readers_guard.Release(); return merge_tree_writer_->WriteSortedReaders(std::move(sorted_readers)); } diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index 6885dc374..34c6ef850 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -77,6 +77,13 @@ Result> AppendOnlyTableRead::CreateReader( std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::shared_ptr& split : splits) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); @@ -92,8 +99,6 @@ Result> AppendOnlyTableRead::CreateReader( } } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (!realtime_splits.empty()) { const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { @@ -106,6 +111,9 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + cleanup_guard.Release(); return result; } @@ -124,6 +132,13 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); std::vector> readers; readers.reserve(realtime_split->DiskSplits().size() + 1); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), realtime_split->Bucket()); if (memory.partition_bucket != expected_partition_bucket) { @@ -150,8 +165,17 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( std::vector> memory_readers, memory.store->CreateQueryReaders(memory.read_view, realtime_split->CommittedEndOffset(), query_context)); - + const size_t first_memory_reader = readers.size(); + readers.reserve(readers.size() + memory_readers.size()); for (std::unique_ptr& memory_reader : memory_readers) { + readers.push_back(std::move(memory_reader)); + } + + for (size_t i = first_memory_reader; i < readers.size(); ++i) { + std::unique_ptr& memory_reader = readers[i]; + if (!memory_reader) { + return Status::Invalid("append-only real-time store returned a null query reader"); + } if (context_->EnablePredicateFilter() && context_->GetPredicate()) { PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create( std::move(memory_reader), @@ -159,14 +183,15 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, RealtimeReader::Create(memory.read_view, std::move(memory_reader))); - readers.push_back(std::move(realtime_reader)); + memory_reader = std::move(realtime_reader); } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (release_ticket) { PAIMON_RETURN_NOT_OK( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + readers_guard.Release(); return result; } diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 64d722097..96f3f00d9 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -209,6 +209,13 @@ Result> KeyValueTableRead::CreateReader( std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::shared_ptr& split : splits) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); @@ -223,8 +230,6 @@ Result> KeyValueTableRead::CreateReader( } } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (!realtime_splits.empty()) { const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { @@ -237,6 +242,9 @@ Result> KeyValueTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + cleanup_guard.Release(); return result; } diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index 1b496d8a1..4c3968dc3 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -152,16 +152,17 @@ Result>> RealtimeTableScan::CreateRealtimeSpl continue; } - // Append tables can schedule all but the tail disk split independently. The tail split - // carries the immutable memory view so disk and memory are still concatenated by one - // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit. - auto tail_disk_split = std::prev(grouped_disk_splits.end()); - result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split); - std::vector> realtime_disk_splits; - realtime_disk_splits.push_back(std::move(*tail_disk_split)); RealtimePartitionBucketView& memory = memory_iter->second; + if (!pk_table_) { + // Append tables can schedule all but the tail disk split independently. The tail split + // carries the immutable memory view so disk and memory are still concatenated by one + // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit. + auto tail_disk_split = std::prev(grouped_disk_splits.end()); + result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split); + grouped_disk_splits.erase(grouped_disk_splits.begin(), tail_disk_split); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_split, - create_realtime_split(key, std::move(realtime_disk_splits), memory)); + create_realtime_split(key, std::move(grouped_disk_splits), memory)); result.push_back(std::move(realtime_split)); active_memory.erase(memory_iter); } @@ -176,7 +177,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl return result; } -RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, +RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -184,6 +185,7 @@ RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, const std::shared_ptr& scan_filter, int64_t read_view_ttl_millis) : disk_scan_(std::move(disk_scan)), + pk_table_(pk_table), realtime_context_(realtime_context), path_factory_(path_factory), snapshot_manager_(snapshot_manager), diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h index 959203ca4..692b749ef 100644 --- a/src/paimon/core/table/source/realtime_table_scan.h +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -38,7 +38,7 @@ class SnapshotManager; /// Adds process-local memory splits to a normal data-table batch scan. class RealtimeTableScan : public TableScan { public: - RealtimeTableScan(std::unique_ptr&& disk_scan, + RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -67,6 +67,7 @@ class RealtimeTableScan : public TableScan { const std::optional& snapshot_id) const; std::unique_ptr disk_scan_; + bool pk_table_; std::shared_ptr realtime_context_; std::shared_ptr path_factory_; std::shared_ptr snapshot_manager_; diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 0bcd79f61..f894e1a74 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -344,7 +344,7 @@ Result> NewDataTableScan(const std::shared_ptr realtime_context, RealtimeContextImpl::Cast(context->GetRealtimeContext())); return std::make_unique( - std::move(batch_scan), realtime_context, path_factory, + std::move(batch_scan), pk_table, realtime_context, path_factory, snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(), context->GetScanFilters(), core_options.GetRealtimeReadViewTtlMillis()); } diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 0e6f83b70..a393f838d 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -468,21 +468,27 @@ class SwapFirstTwoBatchReader final : public BatchReader { Result NextBatch() override { if (!initialized_) { initialized_ = true; - PAIMON_ASSIGN_OR_RAISE(ReadBatch first, delegate_->NextBatch()); - if (BatchReader::IsEofBatch(first)) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { return MakeEofBatch(); } - PAIMON_ASSIGN_OR_RAISE(ReadBatch second, delegate_->NextBatch()); - if (BatchReader::IsEofBatch(second)) { - return first; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (array->length() < 2) { + return Status::Invalid("cannot make a one-row reader unsorted"); + } + arrow::ArrayVector pieces = {array->Slice(1, 1), array->Slice(0, 1)}; + if (array->length() > 2) { + pieces.push_back(array->Slice(2)); } - first_ = std::move(first); - return second; - } - if (first_.has_value()) { - ReadBatch first = std::move(first_.value()); - first_.reset(); - return first; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr swapped, + arrow::Concatenate(pieces)); + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*swapped, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); } return delegate_->NextBatch(); } @@ -492,14 +498,12 @@ class SwapFirstTwoBatchReader final : public BatchReader { } void Close() override { - first_.reset(); delegate_->Close(); } private: bool initialized_ = false; std::unique_ptr delegate_; - std::optional first_; }; class SubstituteOffsetBatchReader final : public BatchReader { @@ -1253,6 +1257,8 @@ class RealtimeWriteInteTest : public ::testing::Test { return Status::OK(); } + void RunConcurrencyTest(bool primary_key); + Result ReadCommittedOffsets() const { PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); @@ -1326,8 +1332,10 @@ class RealtimeWriteInteTest : public ::testing::Test { } else { CreateTable(/*partition_keys=*/{"pt"}); } + auto close_state = std::make_shared(); + auto factory = std::make_shared(close_state); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); + RealtimeContext::Create(factory)); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter(realtime_context)); std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); @@ -1363,6 +1371,9 @@ class RealtimeWriteInteTest : public ::testing::Test { TableRead::Create(std::move(read_context))); ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), "unsupported real-time split version"); + if (!primary_key) { + ASSERT_EQ(1, close_state->query_close_count->load(std::memory_order_acquire)); + } std::vector expected_rows = p0_rows; expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); @@ -1631,6 +1642,56 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkMergeAllDiskSplitsWithMemory) { + options_[Options::SOURCE_SPLIT_OPEN_FILE_COST] = "1"; + options_[Options::SOURCE_SPLIT_TARGET_SIZE] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + {{20, "disk-20", "p0"}, {21, "disk-21", "p0"}}, + }; + for (int64_t commit_identifier = 0; + commit_identifier < static_cast(disk_batches.size()); ++commit_identifier) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_batches[commit_identifier], /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr memory_batch, + MakeBatch({Row{1, "memory-1", "p0"}, Row{10, "deleted-10", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, realtime_split); + ASSERT_EQ(3, realtime_split->DiskSplits().size()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "memory-1", "p0"}, + {2, "disk-2", "p0"}, + {11, "disk-11", "p0"}, + {20, "disk-20", "p0"}, + {21, "disk-21", "p0"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { const std::shared_ptr projected_b = arrow::field("b", arrow::int64()); fields_ = { @@ -1712,6 +1773,45 @@ TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkMemoryReadAfterSchemaEvolution) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "old", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + std::shared_ptr renamed_payload = arrow::field("renamed_payload", arrow::utf8()); + std::shared_ptr added = arrow::field("added", arrow::int32()); + ASSERT_OK(TestHelper::WriteNextSchema(dir_->GetFileSystem(), table_path_, + {DataField(0, fields_[0]), DataField(1, renamed_payload), + DataField(2, fields_[2]), DataField(3, added)}, + /*highest_field_id=*/3, options_)); + fields_[1] = renamed_payload; + fields_.push_back(added); + schema_ = arrow::schema(fields_); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "renamed_payload", "pt", "added"}, + /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + ASSERT_EQ(1, result.data->num_chunks()); + std::shared_ptr row = + std::dynamic_pointer_cast(result.data->chunk(0)); + ASSERT_NE(nullptr, row); + ASSERT_EQ(1, row->length()); + std::shared_ptr renamed_values = + std::dynamic_pointer_cast(row->field(2)); + ASSERT_NE(nullptr, renamed_values); + ASSERT_EQ("old", renamed_values->GetString(0)); + ASSERT_TRUE(row->field(4)->IsNull(0)); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -2277,6 +2377,40 @@ TEST_F(RealtimeWriteInteTest, TestPkQueryReaderCloseFailure) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestAppendQueryReaderCloseFailure) { + CreateTable(/*partition_keys=*/{}); + auto state = std::make_shared(); + state->query_null_index = 1; + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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_NOK_WITH_MSG(table_read->CreateReader(plan->Splits()), + "append-only real-time store returned a null query reader"); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); + + state->query_null_index = -1; + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkCommitReaderCloseFailure) { CreatePkTable(); auto state = std::make_shared(); @@ -3705,8 +3839,12 @@ TEST_F(RealtimeWriteInteTest, TestReopenRealtimeContextAfterRollback) { ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { - CreateTable(/*partition_keys=*/{}); +void RealtimeWriteInteTest::RunConcurrencyTest(bool primary_key) { + if (primary_key) { + CreatePkTable(); + } else { + CreateTable(/*partition_keys=*/{}); + } ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -3942,6 +4080,14 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { + RunConcurrencyTest(/*primary_key=*/false); +} + +TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { + RunConcurrencyTest(/*primary_key=*/true); +} + TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { CreateTable(/*partition_keys=*/{"pt"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, From 75f959a532f5d1f76d786716ba5c1af9ef4eaed7 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Tue, 25 Aug 2026 00:31:55 +0800 Subject: [PATCH 23/24] refactor(realtime): simplify primary-key contracts --- .../realtime/arrow_realtime_store_factory.h | 1 - include/paimon/realtime/realtime_store.h | 39 ++++++++----------- src/paimon/core/mergetree/merge_tree_writer.h | 5 +-- .../realtime/primary_key_realtime_store.h | 2 +- .../realtime/realtime_primary_key_writer.h | 1 - 5 files changed, 19 insertions(+), 29 deletions(-) diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index da1b8de36..153d524d4 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -26,7 +26,6 @@ namespace paimon { /// Factory for Paimon's default Arrow-backed `RealtimeStore`. class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: - /// Creates the built-in append or in-memory primary-key store. Result> Create(RealtimeStoreCreateRequest&& request) override; }; diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index 90c6ce0a8..60d1afc39 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -69,10 +69,10 @@ struct PAIMON_EXPORT RealtimeStoreCreateRequest { /// A record batch and its framework-assigned contiguous offset range. /// -/// Append-mode batches contain table write fields, and row `i` is associated with -/// `offset_range.begin + i`. Primary-key batches contain the prepared transport schema supplied -/// to the factory and are physically sorted by full primary key then sequence number; their -/// per-row `_REALTIME_OFFSET` field preserves the original write-order offset after sorting. +/// Append-mode batches contain table write fields, and row `i` has offset +/// `offset_range.begin + i`. Primary-key batches use the prepared transport schema, are sorted +/// by full primary key then sequence number, and retain the original offset in +/// `_REALTIME_OFFSET`. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -147,14 +147,10 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// - /// Concatenating the returned readers must produce every sealed row exactly once. Append-mode - /// readers preserve write order and contain `_VALUE_KIND` followed by table write fields. - /// Primary-key readers expose raw prepared rows. Each returned reader's complete stream, - /// including across `NextBatch` boundaries, is sorted by full primary key then sequence - /// number; all readers collectively cover sealed mutations exactly once. Reader cardinality is - /// independent of the number of writes. Paimon adapts and merges those rows before writing - /// files. Paimon validates the complete ordering and coverage before publishing generated file - /// state; a violation fails the prepare operation. + /// The returned readers collectively expose every sealed row exactly once. Append-mode readers + /// preserve write order and contain `_VALUE_KIND` followed by table write fields. Primary-key + /// readers use the prepared transport schema; each reader's complete stream is sorted by full + /// primary key then sequence number. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -164,18 +160,15 @@ class PAIMON_EXPORT RealtimeStore { /// also provide a consistent snapshot when a write or seal is in progress. virtual Result> AcquireReadView() = 0; - /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater - /// than or equal to `offset_begin`. Primary-key mode ignores `offset_begin` and returns raw - /// prepared rows; Paimon applies offset filtering, projection, and merge-on-read adaptation. + /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater than + /// or equal to `offset_begin`; primary-key mode ignores `offset_begin`. /// - /// Append-mode output batches contain `_VALUE_KIND` first, followed by requested fields except - /// a duplicate `_VALUE_KIND`; all returned append readers collectively cover every matching - /// row exactly once. Primary-key output batches use the prepared transport schema and may - /// contain multiple mutations per key. Each returned primary-key reader's complete stream is - /// sorted by full primary key then sequence number, and all readers collectively cover raw - /// mutations exactly once. Reader cardinality is independent of the number of writes. Paimon - /// validates ordering while adapting each complete reader stream and retains `view` for the - /// lifetime of the resulting framework reader. + /// Append-mode batches contain `_VALUE_KIND` followed by the requested fields except a duplicate + /// `_VALUE_KIND`, and collectively expose every matching row exactly once. Primary-key batches + /// use the prepared transport schema and may contain multiple mutations per key; each reader's + /// complete stream is sorted by full primary key then sequence number, and the readers + /// collectively expose every raw mutation exactly once. Paimon retains `view` for the lifetime + /// of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index cea07f3e4..01efd975c 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -70,9 +70,8 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; - /// Consumes readers whose complete streams are individually sorted by primary key and - /// sequence number. Readers are closed on success or failure; an error may leave generated - /// file state unpublished, so the caller must discard this writer and replay its input. + /// Consumes readers whose complete streams are individually sorted by primary key and sequence + /// number. Readers are closed on success or failure. Status WriteSortedReaders(std::vector>&& readers); Status Compact(bool full_compaction) override; diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h index d6a23ccf9..f779b4d7d 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store.h +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -35,7 +35,7 @@ class TableSchema; Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema); -/// In-memory store for prepared primary-key real-time batches. +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. class PrimaryKeyRealtimeStore final : public RealtimeStore { public: static Result> Create( diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index 2eaf7ce24..d65c7e533 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -42,7 +42,6 @@ class FieldsComparator; class RealtimeContextImpl; struct RealtimeStoreState; -/// Coordinates framework-prepared primary-key real-time writes. class RealtimePrimaryKeyWriter final : public BatchWriter { public: static Result> Create( From 3f0efbae2ae99058749f30e3cccbdc2940577e47 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Tue, 25 Aug 2026 00:32:24 +0800 Subject: [PATCH 24/24] fix(realtime): strengthen primary-key recovery coverage --- .../core/mergetree/merge_tree_writer_test.cpp | 58 +++++++ .../core/realtime/realtime_context_impl.cpp | 16 +- .../core/realtime/realtime_context_test.cpp | 6 +- .../table/source/key_value_table_read.cpp | 3 + test/inte/realtime_write_inte_test.cpp | 158 ++++++++++++++++++ 5 files changed, 237 insertions(+), 4 deletions(-) diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 63e896573..9ce5498cb 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -530,6 +530,64 @@ TEST_P(MergeTreeWriterTest, TestSortedReaders) { ASSERT_EQ(1, new_file->delete_row_count); } +TEST_P(MergeTreeWriterTest, TestMergeSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/4, dir->Str(), path_factory, + /*schema_id=*/7, options)); + + auto first_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [5, 0, "Alice", 10, 0, 15.1], + [7, 0, "Carol", 20, 1, 17.1], + [10, 0, "Eve", 30, 2, 20.1] + ])") + .ValueOrDie()); + auto second_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1] + ])") + .ValueOrDie()); + bool first_closed = false; + bool second_closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(first_array), &first_closed)); + sorted_readers.push_back(std::make_unique( + CreateSingleReader(second_array), &second_closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(first_closed); + ASSERT_TRUE(second_closed); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(5, new_file->row_count); + ASSERT_EQ(1, new_file->delete_row_count); + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [5, 0, "Alice", 10, 0, 15.1], + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1], + [10, 0, "Eve", 30, 2, 20.1] + ])"}, + &expected_array) + .ok()); + CheckFileContent(path_factory->ToPath(new_file), expected_array); +} + TEST_P(MergeTreeWriterTest, TestSortedReaderOwnership) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index ba4c8b7a6..736ebb02d 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -59,6 +59,17 @@ bool SameMode(const RealtimeStoreCreateConfig& left, const RealtimeStoreCreateCo return true; } +std::string PartitionToString(const std::map& partition) { + std::string result = "{"; + for (auto iter = partition.begin(); iter != partition.end(); ++iter) { + if (iter != partition.begin()) { + result += ", "; + } + result += iter->first + "=" + iter->second; + } + return result + "}"; +} + } // namespace RealtimeContextImpl::RealtimeContextImpl(const std::shared_ptr& factory) @@ -120,8 +131,9 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( if (iter != stores_.end()) { if (!SameMode(iter->second.mode_config, request.mode_config) || !iter->second.write_schema->Equals(*requested_schema, /*check_metadata=*/true)) { - return Status::Invalid( - "real-time store schema or mode does not match the registered store"); + return Status::Invalid("real-time store schema or mode mismatch for partition " + + PartitionToString(key.partition) + ", bucket " + + std::to_string(key.bucket) + "; recreate the RealtimeContext"); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, iter->second.store->AcquireReadView()); diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 2b47e9dc9..916b46aad 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -171,13 +171,15 @@ TEST(RealtimeContextTest, TestRejectsMismatchedSchemaOnStoreReuse) { ASSERT_NOK_WITH_MSG( GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(arrow::int32(), metadata), {}, GetDefaultPool()), - "schema or mode does not match"); + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); ASSERT_NOK_WITH_MSG( GetOrCreateAppendStore( context, partition, 0, MakeWriteSchema(arrow::int64(), arrow::key_value_metadata({"identity"}, {"v2"})), {}, GetDefaultPool()), - "schema or mode does not match"); + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); ASSERT_EQ(1, factory->stores.size()); } diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 96f3f00d9..3532b59b4 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -253,6 +253,9 @@ Result> KeyValueTableRead::CreateRealtimeReader( if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { return Status::Invalid("unsupported real-time split version"); } + if (realtime_split->MemoryEndOffset() < realtime_split->CommittedEndOffset()) { + return Status::Invalid("real-time split memory end offset precedes committed end offset"); + } const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { return Status::Invalid("reading a real-time split requires a real-time context"); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index a393f838d..53737fc2b 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -1281,6 +1281,30 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_EQ(expected_rows, actual_rows); } + void ReplayPkWalAndCommit(const std::vector& wal, + const std::vector& row_kinds, + int64_t commit_identifier, + const std::vector& expected_rows) const { + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + } + void CheckDropDatePartitionRemovesOffset(bool legacy_partition_name_enabled) { fields_ = {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::date32())}; @@ -2861,6 +2885,44 @@ TEST_F(RealtimeWriteInteTest, TestPlanExcludesRowsWrittenAfterMemoryEndOffset) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkRejectsReversedVisibleOffsets) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + std::vector> disk_splits = split->DiskSplits(); + std::vector> invalid_splits = {std::make_shared( + split->Version(), split->SnapshotId(), split->Partition(), split->Bucket(), + std::move(disk_splits), split->MemoryEndOffset() + 1, split->MemoryEndOffset(), + split->OpaqueTicket())}; + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + 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_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), + "memory end offset precedes committed end offset"); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestReadFailsAfterRealtimeSplitTicketExpires) { options_[Options::REALTIME_READ_VIEW_TTL] = "10 ms"; CreateTable(/*partition_keys=*/{}); @@ -4392,4 +4454,100 @@ TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { ASSERT_EQ(5, second_committed_offsets.at(partition_bucket)); } +TEST_F(RealtimeWriteInteTest, TestPkExternalCommitRecovery) { + CreatePkTable(); + const std::vector seed_rows = {{99, "seed", "p0"}}; + ReplayPkWalAndCommit(seed_rows, /*row_kinds=*/{}, /*commit_identifier=*/0, seed_rows); + + const std::vector wal = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector row_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector failed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, failed_progress.size()); + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result failed_commit = + commit->CommitWithProgress(failed_progress, /*commit_identifier=*/1, + /*watermark=*/std::nullopt); + io_hook->Clear(); + ASSERT_TRUE(failed_commit.status().IsIOError()) << failed_commit.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(seed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}; + ReplayPkWalAndCommit(wal, row_kinds, /*commit_identifier=*/1, expected_rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkRefreshRecovery) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + + const std::vector base_rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(base_rows, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t base_snapshot_id, Commit(base_progress, /*commit_identifier=*/0)); + ASSERT_OK(failed_writer->RefreshCommittedSnapshot(base_snapshot_id)); + + const std::vector committed_wal = { + {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector committed_kinds = {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr committed_batch, + MakeBatch(committed_wal, /*partitioned=*/false, /*bucket=*/0, committed_kinds)); + ASSERT_OK(failed_writer->Write(std::move(committed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector committed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id, + Commit(committed_progress, /*commit_identifier=*/1)); + + const std::vector replay_wal = {{4, "four", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(replay_wal, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(replay_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Status failed_refresh = failed_writer->RefreshCommittedSnapshot(committed_snapshot_id); + io_hook->Clear(); + ASSERT_TRUE(failed_refresh.IsIOError()) << failed_refresh.ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + const std::vector committed_rows = {{1, "one-new", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(committed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}; + ReplayPkWalAndCommit(replay_wal, /*row_kinds=*/{}, /*commit_identifier=*/2, expected_rows); +} + } // namespace paimon::test