diff --git a/include/paimon/file_store_write.h b/include/paimon/file_store_write.h index 2c4e87086..3e3cfb8e4 100644 --- a/include/paimon/file_store_write.h +++ b/include/paimon/file_store_write.h @@ -60,6 +60,15 @@ class PAIMON_EXPORT FileStoreWrite { /// the corresponding array in `batch` must have zero null entries. virtual Status Write(std::unique_ptr&& batch) = 0; + /// Slices the current in-memory real-time data into a sealed segment so it can be reclaimed + /// independently. A future implementation will support spilling sealed segments to a + /// temporary directory; currently this method only creates the in-memory segment boundary. + /// Calling this method on a non-real-time writer returns an error. + /// If sealing fails, the caller must recreate both the `RealtimeContext` and writer. The + /// upstream must then recover input from the durable recovery offset persisted in the + /// snapshot. Reusing the failed writer is unsupported. + virtual Status Seal(); + /// Compact data stored in given partition and bucket. Note that compaction process is only /// submitted and may not be completed when the method returns. /// @@ -99,13 +108,17 @@ class PAIMON_EXPORT FileStoreWrite { /// Generates commit messages together with partition-bucket real-time offset ranges. /// /// Each range is returned atomically with the commit message generated from the same sealed - /// segment. Repeated calls return incremental progress. The upstream coordinator must retain + /// segments. Repeated calls return incremental progress. The upstream coordinator must retain /// every result until it is committed and include all earlier prepared-but-uncommitted /// progress when a later checkpoint subsumes it. /// /// @param commit_identifier Identifier of this prepare-commit operation in streaming mode. /// @return Real-time commit messages with their partition-bucket offset ranges. /// @note Calling this method on a non-real-time writer or in batch mode returns an error. + /// @note If preparation fails, the caller must recreate both the `RealtimeContext` and writer. + /// The upstream must then recover input from the durable recovery offset persisted in the + /// snapshot. The failed writer may contain partially prepared bucket state and must not + /// be reused. virtual Result> PrepareCommitWithProgress( int64_t commit_identifier); @@ -122,6 +135,12 @@ class PAIMON_EXPORT FileStoreWrite { virtual Status RefreshCommittedSnapshot(int64_t snapshot_id); virtual std::shared_ptr GetMetrics() const = 0; + + /// Releases resources owned by this writer. + /// + /// Closing a real-time writer with data not covered by a successful + /// `PrepareCommitWithProgress()` returns an error and invalidates its `RealtimeContext`. + /// The caller must rebuild both objects and recover input from the durable snapshot offset. virtual Status Close() = 0; }; diff --git a/include/paimon/realtime/realtime_commit_progress.h b/include/paimon/realtime/realtime_commit_progress.h index a25fdbf5c..bf19695c3 100644 --- a/include/paimon/realtime/realtime_commit_progress.h +++ b/include/paimon/realtime/realtime_commit_progress.h @@ -35,9 +35,9 @@ namespace paimon { /// count is not the represented row count. The progress fields are not embedded in `CommitMessage` /// serialization. struct PAIMON_EXPORT RealtimeCommitProgress { - /// Paimon commit message generated from one sealed segment. + /// Paimon commit message generated from one or more sealed segments. std::shared_ptr commit_message; - /// Partition-bucket containing the sealed segment. + /// Partition-bucket containing the sealed segments. RealtimePartitionBucket partition_bucket; /// Left-closed, right-open offset range represented by the commit message. OffsetRange offset_range; diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index c4e108dbf..115eddce0 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -50,9 +50,9 @@ enum class PAIMON_EXPORT RealtimeStoreMode { /// Parameters used by a `RealtimeStoreFactory` to create a store. struct PAIMON_EXPORT RealtimeStoreCreateRequest { /// Schema whose ownership is transferred to the factory. Append mode receives the complete - /// append transport schema: [_REALTIME_OFFSET, table write fields]. Primary-key mode receives - /// the realtime primary-key transport schema: - /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields]. + /// append store-write schema: [_REALTIME_OFFSET, table write fields]. Primary-key mode receives + /// the real-time primary-key store-write schema: + /// [_SEQUENCE_NUMBER, _VALUE_KIND, _REALTIME_OFFSET, table write fields]. std::unique_ptr<::ArrowSchema> write_schema; /// Table options available to the store implementation. std::map options; @@ -66,12 +66,13 @@ struct PAIMON_EXPORT RealtimeStoreCreateRequest { /// A record batch and its application-assigned offset bounds. /// -/// Append-mode batches use the append transport schema [_REALTIME_OFFSET, table write fields], and -/// offsets are strictly increasing before the batch enters the store. Primary-key batches use the -/// realtime primary-key transport schema, are sorted by full primary key then sequence number, and -/// retain the original offset in `_REALTIME_OFFSET`. `offset_range` is the left-closed, right-open -/// envelope from the first application offset through one past the last; offsets may have gaps, so -/// its count is not the batch row count. +/// Append-mode batches use the append store-write schema +/// [_REALTIME_OFFSET, table write fields], and offsets are strictly increasing before the batch +/// enters the store. Primary-key batches use the real-time primary-key store-write schema, are +/// sorted by full primary key then sequence number, and retain the original offset in +/// `_REALTIME_OFFSET`. `offset_range` is the left-closed, right-open envelope from the first +/// application offset through one past the last; offsets may have gaps, so its count is not the +/// batch row count. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -109,11 +110,11 @@ class PAIMON_EXPORT RealtimeReadView { /// Parameters used by a `RealtimeStore` to create readers for a query. struct PAIMON_EXPORT RealtimeQueryContext { - /// Physical source schema the store must materialize. Query readers must include the mandatory - /// `_VALUE_KIND` field in returned batches. Paimon may subsequently convert physical fields - /// into the query's logical output schema, for example for selected-key MAP or VARIANT access. - /// This schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must - /// import or copy it synchronously. + /// Physical source schema the store must materialize. Every returned batch must match this + /// schema exactly. Paimon may subsequently add framework fields or convert physical fields into + /// the query's logical output schema, for example for selected-key MAP or VARIANT access. This + /// schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must import + /// or copy it synchronously. ::ArrowSchema* read_schema; /// Optional predicate using field indexes from `read_schema`. A non-null predicate allows the /// plugin to prune candidate rows. Exact filtering is applied by the Paimon read framework. @@ -146,9 +147,9 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// /// The returned readers collectively expose every sealed row exactly once. Append-mode readers - /// preserve write order and contain `_VALUE_KIND`, `_REALTIME_OFFSET`, and table write fields. - /// Primary-key readers contain the realtime primary-key transport fields; each reader's - /// complete stream is sorted by full primary key then sequence number. + /// preserve write order and contain `_REALTIME_OFFSET` followed by the table write fields. + /// Primary-key readers contain the real-time primary-key store fields; each reader's complete + /// stream is sorted by full primary key then sequence number. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 4c9b9b348..01ba7b8bc 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -389,10 +389,10 @@ set(PAIMON_CORE_SRCS core/realtime/arrow_realtime_store_factory.cpp core/realtime/realtime_offset_batch_reader.cpp core/realtime/realtime_primary_key_reader.cpp - core/realtime/primary_key_realtime_store.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp + core/realtime/realtime_schema_layout.cpp core/realtime/realtime_store_read_pipeline.cpp core/realtime/realtime_primary_key_writer.cpp core/postpone/postpone_bucket_writer.cpp @@ -804,6 +804,7 @@ if(PAIMON_BUILD_TESTS) core/realtime/primary_key_realtime_store_test.cpp core/realtime/realtime_offset_batch_reader_test.cpp core/realtime/realtime_store_read_pipeline_test.cpp + core/realtime/realtime_schema_layout_test.cpp core/realtime/realtime_primary_key_reader_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp index c9591df8f..7b6a0985f 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader.cpp +++ b/src/paimon/common/reader/complete_row_kind_batch_reader.cpp @@ -72,13 +72,17 @@ Result CompleteRowKindBatchReader::NextBatchWi PAIMON_ASSIGN_OR_RAISE(std::shared_ptr row_kind_array, PrepareRowKindArray(struct_array->length())); // complete row kind - UpdateFieldNamesWithRowKind(struct_array); arrow::ArrayVector fields_with_row_kind = {row_kind_array}; fields_with_row_kind.insert(fields_with_row_kind.end(), struct_array->fields().begin(), struct_array->fields().end()); + arrow::FieldVector schema_fields_with_row_kind = { + arrow::field(SpecialFields::ValueKind().Name(), arrow::int8())}; + const arrow::FieldVector& schema_fields = struct_array->struct_type()->fields(); + schema_fields_with_row_kind.insert(schema_fields_with_row_kind.end(), schema_fields.begin(), + schema_fields.end()); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr array_with_row_kind, - arrow::StructArray::Make(fields_with_row_kind, field_names_with_row_kind_)); + arrow::StructArray::Make(fields_with_row_kind, schema_fields_with_row_kind)); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*array_with_row_kind, c_array.get(), c_schema.get())); PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); @@ -99,19 +103,4 @@ Result> CompleteRowKindBatchReader::PrepareRowKind } } -void CompleteRowKindBatchReader::UpdateFieldNamesWithRowKind( - const std::shared_ptr& struct_array) { - if (static_cast(struct_array->struct_type()->num_fields()) + 1 == - field_names_with_row_kind_.size()) { - return; - } - field_names_with_row_kind_.clear(); - const auto& fields = struct_array->struct_type()->fields(); - field_names_with_row_kind_.reserve(fields.size() + 1); - field_names_with_row_kind_.push_back(SpecialFields::ValueKind().Name()); - for (const auto& field : fields) { - field_names_with_row_kind_.push_back(field->name()); - } -} - } // namespace paimon diff --git a/src/paimon/common/reader/complete_row_kind_batch_reader.h b/src/paimon/common/reader/complete_row_kind_batch_reader.h index e2581aba3..358bc34fa 100644 --- a/src/paimon/common/reader/complete_row_kind_batch_reader.h +++ b/src/paimon/common/reader/complete_row_kind_batch_reader.h @@ -21,9 +21,7 @@ #include #include -#include #include -#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -52,7 +50,6 @@ class CompleteRowKindBatchReader : public BatchReader { void Close() override { reader_->Close(); row_kind_array_.reset(); - field_names_with_row_kind_.clear(); } std::shared_ptr GetReaderMetrics() const override { @@ -62,12 +59,9 @@ class CompleteRowKindBatchReader : public BatchReader { private: Result> PrepareRowKindArray(int32_t struct_array_length); - void UpdateFieldNamesWithRowKind(const std::shared_ptr& struct_array); - private: std::shared_ptr arrow_pool_; std::unique_ptr reader_; std::shared_ptr row_kind_array_; - std::vector field_names_with_row_kind_; }; } // namespace paimon diff --git a/src/paimon/core/append/append_compact_coordinator.cpp b/src/paimon/core/append/append_compact_coordinator.cpp index b6afc15c4..06716bac7 100644 --- a/src/paimon/core/append/append_compact_coordinator.cpp +++ b/src/paimon/core/append/append_compact_coordinator.cpp @@ -165,7 +165,8 @@ std::unique_ptr CreateFileStoreWrite( path_factory, snapshot_manager, schema_manager, /*commit_user=*/"compact-coordinator", /*root_path=*/table_path, table_schema, arrow_schema, - /*write_schema=*/arrow_schema, partition_schema, + /*write_schema=*/arrow_schema, + /*realtime_schema_layout=*/nullptr, partition_schema, /*dv_maintainer_factory=*/nullptr, /*io_manager=*/nullptr, core_options, /*ignore_previous_files=*/true, diff --git a/src/paimon/core/operation/abstract_file_store_write.cpp b/src/paimon/core/operation/abstract_file_store_write.cpp index e005619a8..c4da29bf5 100644 --- a/src/paimon/core/operation/abstract_file_store_write.cpp +++ b/src/paimon/core/operation/abstract_file_store_write.cpp @@ -31,7 +31,8 @@ #include "paimon/core/operation/file_system_write_restore.h" #include "paimon/core/operation/metrics/compaction_metrics.h" #include "paimon/core/operation/restore_files.h" -#include "paimon/core/realtime/realtime_offset_utils.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" #include "paimon/core/table/bucket_mode.h" @@ -59,6 +60,7 @@ AbstractFileStoreWrite::AbstractFileStoreWrite( const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, @@ -73,6 +75,7 @@ AbstractFileStoreWrite::AbstractFileStoreWrite( root_path_(root_path), schema_(schema), write_schema_(write_schema), + realtime_schema_layout_(realtime_schema_layout), table_schema_(table_schema), partition_schema_(partition_schema), dv_maintainer_factory_(dv_maintainer_factory), @@ -95,6 +98,7 @@ Status AbstractFileStoreWrite::Write(std::unique_ptr&& batch) { if (PAIMON_UNLIKELY(batch == nullptr)) { return Status::Invalid("batch is null pointer"); } + PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable()); // in FileStoreWrite::Create() we have checked the table kind and bucket mode, here we only // check the bucket id in batch if (options_.GetBucket() == -1) { @@ -123,8 +127,11 @@ Status AbstractFileStoreWrite::Write(std::unique_ptr&& batch) { } } // check nullability - std::shared_ptr input_schema = - IsRealtimeWrite() ? RealtimeOffsetUtils::CreateInputSchema(write_schema_) : write_schema_; + if (IsRealtimeWrite() && !realtime_schema_layout_) { + return Status::Invalid("real-time schema layout is null"); + } + const std::shared_ptr& input_schema = + realtime_schema_layout_ ? realtime_schema_layout_->InputSchema() : write_schema_; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr data, arrow::ImportArray(batch->GetData(), arrow::struct_(input_schema->fields()))); @@ -140,8 +147,34 @@ Status AbstractFileStoreWrite::Write(std::unique_ptr&& batch) { return writer_memory_manager_->OnWriteCompleted(writer.get()); } +Status AbstractFileStoreWrite::Seal() { + if (!IsRealtimeWrite()) { + return FileStoreWrite::Seal(); + } + std::lock_guard seal_prepare_lock(realtime_seal_prepare_mutex_); + PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable()); + std::vector> writers; + { + std::lock_guard lock(writers_mutex_); + for (const auto& [_, bucket_writers] : writers_) { + for (const auto& [_, writer_container] : bucket_writers) { + writers.push_back(writer_container.writer); + } + } + } + for (const std::shared_ptr& writer : writers) { + Status status = writer->Seal(); + if (!status.ok()) { + // Earlier buckets may already be sealed, so the writer cannot safely continue. + return PoisonRealtimeWrite(status); + } + } + return Status::OK(); +} + Status AbstractFileStoreWrite::Compact(const std::map& partition, int32_t bucket, bool full_compaction) { + PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable()); PAIMON_ASSIGN_OR_RAISE(BinaryRow part, file_store_path_factory_->ToBinaryRow(partition)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr writer, GetWriter(part, bucket)); assert(writer); @@ -276,10 +309,18 @@ Result> AbstractFileStoreWrite::PrepareCommi if (is_streaming_mode_ == false) { return Status::Invalid("PrepareCommitWithProgress requires streaming mode"); } + // Keep one seal/prepare attempt in flight so a failure fences every later checkpoint. + std::lock_guard seal_prepare_lock(realtime_seal_prepare_mutex_); + PAIMON_RETURN_NOT_OK(CheckRealtimeWriteUsable()); // Real-time prepare snapshots writers under a short lock so writes can continue while sealed // segments are flushed. The normal path iterates and may erase writers in place, which would // either race with concurrent writer creation or hold writers_mutex_ for the whole prepare. - return PrepareRealtimeCommit(); + Result> result = PrepareRealtimeCommit(); + if (!result.ok()) { + // Prepare may have drained an earlier bucket or partially written the failing bucket. + return PoisonRealtimeWrite(result.status()); + } + return result; } Result> AbstractFileStoreWrite::PrepareRealtimeCommit() { @@ -338,22 +379,91 @@ Result> AbstractFileStoreWrite::PrepareRealt } Status AbstractFileStoreWrite::Close() { + std::unique_lock seal_prepare_lock(realtime_seal_prepare_mutex_, std::defer_lock); + if (IsRealtimeWrite()) { + seal_prepare_lock.lock(); + } std::lock_guard lock(writers_mutex_); + bool has_unprepared_realtime_data = false; + Status unprepared_error = Status::OK(); + if (IsRealtimeWrite()) { + for (const auto& [_, bucket_writers] : writers_) { + for (const auto& [_, writer_container] : bucket_writers) { + if (writer_container.writer->HasUnpreparedRealtimeData()) { + has_unprepared_realtime_data = true; + break; + } + } + if (has_unprepared_realtime_data) { + break; + } + } + if (has_unprepared_realtime_data) { + // Close cannot publish pending segments for its caller. + unprepared_error = PoisonRealtimeWrite( + Status::Invalid("real-time writer closed with data not covered by a successful " + "PrepareCommitWithProgress")); + } + } + Status first_error = Status::OK(); for (auto& [_, bucket_writers] : writers_) { for (auto& [_, writer_container] : bucket_writers) { writer_memory_manager_->UnregisterWriter(writer_container.writer.get()); - PAIMON_RETURN_NOT_OK(writer_container.writer->Close()); + Status status = writer_container.writer->Close(); + if (!status.ok() && first_error.ok()) { + first_error = std::move(status); + } } } writers_.clear(); compact_executor_->ShutdownNow(); - return Status::OK(); + if (!first_error.ok() && IsRealtimeWrite()) { + return PoisonRealtimeWrite(first_error); + } + if (has_unprepared_realtime_data) { + return unprepared_error; + } + return first_error; } std::shared_ptr AbstractFileStoreWrite::GetMetrics() const { return metrics_; } +Status AbstractFileStoreWrite::CheckRealtimeWriteUsable() const { + if (IsRealtimeWrite() && realtime_write_poisoned_.load()) { + return Status::Invalid( + "real-time writer cannot be reused after Seal or PrepareCommitWithProgress failed, " + "or after it closed with unprepared data; recreate RealtimeContext and " + "FileStoreWrite, then let upstream recover input from the durable recovery offset " + "persisted in the snapshot"); + } + std::shared_ptr context = GetRealtimeContext(); + if (!context) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context_impl, + RealtimeContextImpl::Cast(context)); + return context_impl->CheckUsable(); +} + +void AbstractFileStoreWrite::InvalidateRealtimeContext() { + std::shared_ptr context = + std::dynamic_pointer_cast(GetRealtimeContext()); + if (context) { + context->Invalidate(); + } +} + +Status AbstractFileStoreWrite::PoisonRealtimeWrite(const Status& cause) { + realtime_write_poisoned_.store(true); + InvalidateRealtimeContext(); + return cause.WithMessage(cause.message(), + "; recreate RealtimeContext and FileStoreWrite, then let upstream " + "recover input from the durable recovery offset persisted in the " + "snapshot"); +} + int32_t AbstractFileStoreWrite::GetDefaultBucketNum() const { return options_.GetBucket(); } diff --git a/src/paimon/core/operation/abstract_file_store_write.h b/src/paimon/core/operation/abstract_file_store_write.h index f43046bfd..25955863f 100644 --- a/src/paimon/core/operation/abstract_file_store_write.h +++ b/src/paimon/core/operation/abstract_file_store_write.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -65,6 +66,8 @@ class MemoryPool; class RecordBatch; class RestoreFiles; class IOManager; +class RealtimeContext; +class RealtimeSchemaLayout; class AbstractFileStoreWrite : public FileStoreWrite { public: @@ -77,6 +80,7 @@ class AbstractFileStoreWrite : public FileStoreWrite { const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, @@ -84,6 +88,7 @@ class AbstractFileStoreWrite : public FileStoreWrite { const std::shared_ptr& executor, const std::shared_ptr& pool); Status Write(std::unique_ptr&& batch) override; + Status Seal() override; Status Compact(const std::map& partition, int32_t bucket, bool full_compaction) override; @@ -123,6 +128,10 @@ class AbstractFileStoreWrite : public FileStoreWrite { return false; } + virtual std::shared_ptr GetRealtimeContext() const { + return nullptr; + } + Result> ScanExistingFileMetas(const BinaryRow& partition, int32_t bucket) const; int32_t GetDefaultBucketNum() const; @@ -136,6 +145,7 @@ class AbstractFileStoreWrite : public FileStoreWrite { std::string root_path_; std::shared_ptr schema_; std::shared_ptr write_schema_; + std::shared_ptr realtime_schema_layout_; std::shared_ptr table_schema_; std::shared_ptr partition_schema_; std::shared_ptr dv_maintainer_factory_; @@ -148,6 +158,10 @@ class AbstractFileStoreWrite : public FileStoreWrite { std::shared_ptr compaction_metrics_; private: + Status CheckRealtimeWriteUsable() const; + void InvalidateRealtimeContext(); + Status PoisonRealtimeWrite(const Status& cause); + Result> GetWriter(const BinaryRow& partition, int32_t bucket); Result> PrepareRealtimeCommit(); @@ -155,11 +169,13 @@ class AbstractFileStoreWrite : public FileStoreWrite { std::unordered_map>> writers_; std::mutex writers_mutex_; + std::mutex realtime_seal_prepare_mutex_; std::mutex realtime_metrics_mutex_; bool ignore_previous_files_ = false; bool is_streaming_mode_ = false; bool ignore_num_bucket_check_ = false; bool batch_committed_ = false; + std::atomic realtime_write_poisoned_{false}; std::shared_ptr metrics_; std::unique_ptr logger_; diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index f91e0f60d..55853de97 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -93,6 +93,7 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, @@ -100,9 +101,10 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool) : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, - root_path, table_schema, schema, write_schema, partition_schema, - dv_maintainer_factory, io_manager, options, ignore_previous_files, - is_streaming_mode, ignore_num_bucket_check, executor, pool), + root_path, table_schema, schema, write_schema, realtime_schema_layout, + partition_schema, dv_maintainer_factory, io_manager, options, + ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, + executor, pool), realtime_context_(realtime_context), logger_(Logger::GetLogger("AppendOnlyFileStoreWrite")) { write_cols_ = write_schema->field_names(); @@ -294,11 +296,8 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( file_store_path_factory_->GeneratePartitionVector(partition)); std::map partition_map(partition_values.begin(), partition_values.end()); - auto c_write_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*write_schema_, c_write_schema.get())); - return RealtimeAppendOnlyWriter::Create( - partition_map, bucket, std::move(c_write_schema), realtime_context_, writer, write_schema_, - options_.GetRealtimeStoreStatisticsMode(), options_.ToMap(), pool_); + return RealtimeAppendOnlyWriter::Create(partition_map, bucket, realtime_context_, writer, + realtime_schema_layout_, options_, pool_); } Result AppendOnlyFileStoreWrite::GetDataFileWriterFactory( diff --git a/src/paimon/core/operation/append_only_file_store_write.h b/src/paimon/core/operation/append_only_file_store_write.h index d030a07cc..9455594b8 100644 --- a/src/paimon/core/operation/append_only_file_store_write.h +++ b/src/paimon/core/operation/append_only_file_store_write.h @@ -79,6 +79,7 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, const std::shared_ptr& write_schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, @@ -119,6 +120,10 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { return realtime_context_ != nullptr; } + std::shared_ptr GetRealtimeContext() const override { + return realtime_context_; + } + /// @param plan_factory The active shredding write plan, or nullptr when the rewrite stays a /// plain passthrough. Resolved by the caller because /// `GetDictionaryPassthroughVetoReason` needs the same answer. diff --git a/src/paimon/core/operation/append_only_file_store_write_test.cpp b/src/paimon/core/operation/append_only_file_store_write_test.cpp index ba37f1343..612f931f0 100644 --- a/src/paimon/core/operation/append_only_file_store_write_test.cpp +++ b/src/paimon/core/operation/append_only_file_store_write_test.cpp @@ -45,7 +45,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_offset_utils.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/snapshot.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/utils/snapshot_manager.h" @@ -278,7 +278,10 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestRealtimeWriteTracksExternalOffsetRange) }; auto logical_schema = arrow::schema({arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())}); - auto realtime_schema = RealtimeOffsetUtils::CreateInputSchema(logical_schema); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, logical_schema)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); auto dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); CreateTable(dir->Str(), logical_schema, options); diff --git a/src/paimon/core/operation/file_store_scan.h b/src/paimon/core/operation/file_store_scan.h index ea963a716..80d7d96bb 100644 --- a/src/paimon/core/operation/file_store_scan.h +++ b/src/paimon/core/operation/file_store_scan.h @@ -132,6 +132,11 @@ class FileStoreScan { return this; } + virtual FileStoreScan* EnableValueFilterForLevels( + const std::function& level_filter) { + return this; + } + /// Drop value statistics from entries after all scan filters have been applied. /// /// @return This scan for chained configuration. diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index fa1294d1b..7f7a74c82 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -37,6 +37,7 @@ #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" #include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/table/bucket_mode.h" @@ -62,14 +63,15 @@ namespace { Status RestoreRealtimeCommittedProgress(const std::shared_ptr& realtime_context, const std::shared_ptr& snapshot_manager, const CoreOptions& options) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_RETURN_NOT_OK(realtime_context_impl->CheckUsable()); 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(realtime_context)); PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( latest_snapshot->Id(), realtime_committed_offsets)); } @@ -82,6 +84,10 @@ Result> FileStoreWrite::PrepareCommitWithPro return Status::Invalid("prepare commit with progress requires a real-time writer"); } +Status FileStoreWrite::Seal() { + return Status::Invalid("seal requires a real-time writer"); +} + Status FileStoreWrite::RefreshCommittedSnapshot(int64_t) { return Status::Invalid("refresh committed snapshot requires a real-time writer"); } @@ -184,6 +190,14 @@ Result> FileStoreWrite::Create(std::unique_ptr realtime_schema_layout; + if (ctx->GetRealtimeContext()) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, write_schema)); + realtime_schema_layout = std::move(layout); + } + std::shared_ptr dv_maintainer_factory; if (need_dv_maintainer_factory) { PAIMON_ASSIGN_OR_RAISE( @@ -201,8 +215,8 @@ Result> FileStoreWrite::Create(std::unique_ptr( file_store_path_factory, snapshot_manager, schema_manager, ctx->GetCommitUser(), - ctx->GetRootPath(), schema, arrow_schema, write_schema, partition_schema, - dv_maintainer_factory, io_manager, options, ignore_previous_files, + ctx->GetRootPath(), schema, arrow_schema, write_schema, realtime_schema_layout, + partition_schema, dv_maintainer_factory, io_manager, options, ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), ctx->GetRealtimeContext(), ctx->GetExecutor(), ctx->GetMemoryPool()); return std::unique_ptr(std::move(file_store_write)); @@ -269,13 +283,21 @@ Result> FileStoreWrite::Create(std::unique_ptr(index_file_handler); } + std::shared_ptr realtime_schema_layout; + if (ctx->GetRealtimeContext()) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, arrow_schema)); + realtime_schema_layout = std::move(layout); + } + return std::make_unique( file_store_path_factory, snapshot_manager, schema_manager, ctx->GetCommitUser(), - ctx->GetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory, - io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options, - ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), - ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), - ctx->GetMemoryPool()); + ctx->GetRootPath(), schema, arrow_schema, realtime_schema_layout, partition_schema, + dv_maintainer_factory, io_manager, key_comparator, sequence_fields_comparator, + merge_function_wrapper, options, ignore_previous_files, ctx->IsStreamingMode(), + ctx->IgnoreNumBucketCheck(), ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), + ctx->GetExecutor(), ctx->GetMemoryPool()); } } diff --git a/src/paimon/core/operation/key_value_file_store_scan.cpp b/src/paimon/core/operation/key_value_file_store_scan.cpp index 54af98a0c..ab3729dfd 100644 --- a/src/paimon/core/operation/key_value_file_store_scan.cpp +++ b/src/paimon/core/operation/key_value_file_store_scan.cpp @@ -76,7 +76,7 @@ Result> KeyValueFileStoreScan::Create( } Result KeyValueFileStoreScan::FilterByStats(const ManifestEntry& entry) const { - PAIMON_ASSIGN_OR_RAISE(bool value_filter_enabled, IsValueFilterEnabled()); + PAIMON_ASSIGN_OR_RAISE(bool value_filter_enabled, IsValueFilterEnabled(entry.Level())); if (value_filter_enabled) { PAIMON_ASSIGN_OR_RAISE(bool filtered, FilterByValueFilter(entry)); if (!filtered) { @@ -168,13 +168,14 @@ Status KeyValueFileStoreScan::SplitAndSetKeyValueFilter( return Status::OK(); } -Result KeyValueFileStoreScan::IsValueFilterEnabled() const { +Result KeyValueFileStoreScan::IsValueFilterEnabled(int32_t level) const { if (value_filter_ == nullptr) { return false; } switch (scan_mode_) { case ScanMode::ALL: - return value_filter_force_enabled_; + return value_filter_force_enabled_ && + (value_filter_level_filter_ == nullptr || value_filter_level_filter_(level)); case ScanMode::DELTA: return false; case ScanMode::CHANGELOG: { diff --git a/src/paimon/core/operation/key_value_file_store_scan.h b/src/paimon/core/operation/key_value_file_store_scan.h index 6438e35d1..f7b02bfaa 100644 --- a/src/paimon/core/operation/key_value_file_store_scan.h +++ b/src/paimon/core/operation/key_value_file_store_scan.h @@ -18,6 +18,8 @@ #pragma once +#include +#include #include #include #include @@ -60,6 +62,14 @@ class KeyValueFileStoreScan : public FileStoreScan { FileStoreScan* EnableValueFilter() override { value_filter_force_enabled_ = true; + value_filter_level_filter_ = nullptr; + return this; + } + + FileStoreScan* EnableValueFilterForLevels( + const std::function& level_filter) override { + value_filter_force_enabled_ = true; + value_filter_level_filter_ = level_filter; return this; } @@ -67,7 +77,8 @@ class KeyValueFileStoreScan : public FileStoreScan { Result FilterByStats(const ManifestEntry& entry) const override; bool WholeBucketFilterEnabled() const override { - return value_filter_ != nullptr && scan_mode_ == ScanMode::ALL; + return value_filter_ != nullptr && scan_mode_ == ScanMode::ALL && + value_filter_level_filter_ == nullptr; } Result> FilterWholeBucketByStats( @@ -86,7 +97,7 @@ class KeyValueFileStoreScan : public FileStoreScan { Status SplitAndSetKeyValueFilter(const std::vector& trimmed_pk); - Result IsValueFilterEnabled() const; + Result IsValueFilterEnabled(int32_t level) const; Result FilterByValueFilter(const ManifestEntry& entry) const; @@ -110,6 +121,7 @@ class KeyValueFileStoreScan : public FileStoreScan { private: bool value_filter_force_enabled_ = false; + std::function value_filter_level_filter_; std::shared_ptr key_filter_; std::shared_ptr value_filter_; std::shared_ptr evolutions_; diff --git a/src/paimon/core/operation/key_value_file_store_scan_test.cpp b/src/paimon/core/operation/key_value_file_store_scan_test.cpp index f6c13cb3e..dd54f290c 100644 --- a/src/paimon/core/operation/key_value_file_store_scan_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_scan_test.cpp @@ -440,6 +440,15 @@ TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithValueStatsCols) { ASSERT_OK_AND_ASSIGN(keep, scan->FilterByStats(entry_keep)); ASSERT_TRUE(keep); + + scan->EnableValueFilterForLevels([](int32_t level) -> bool { return level > 0; }); + ASSERT_FALSE(scan->WholeBucketFilterEnabled()); + entry.File()->level = 0; + ASSERT_OK_AND_ASSIGN(keep, scan->FilterByStats(entry)); + ASSERT_TRUE(keep); + entry.File()->level = 1; + ASSERT_OK_AND_ASSIGN(keep, scan->FilterByStats(entry)); + ASSERT_FALSE(keep); } TEST_F(KeyValueFileStoreScanTest, TestFilterByValueFilterWithSchemaEvolution) { 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 c90d9d05c..cfa62873c 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -35,9 +35,8 @@ #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_offset_utils.h" -#include "paimon/core/realtime/realtime_primary_key_reader.h" #include "paimon/core/realtime/realtime_primary_key_writer.h" +#include "paimon/core/realtime/realtime_schema_layout.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" @@ -62,6 +61,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr& schema_manager, const std::string& commit_user, const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, @@ -74,9 +74,9 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr& executor, const std::shared_ptr& pool) : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, table_schema, schema, /*write_schema=*/schema, - partition_schema, dv_maintainer_factory, io_manager, options, - ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, - executor, pool), + realtime_schema_layout, partition_schema, dv_maintainer_factory, + io_manager, options, ignore_previous_files, is_streaming_mode, + ignore_num_bucket_check, executor, pool), enable_multi_thread_spill_(enable_multi_thread_spill), realtime_context_(realtime_context), key_comparator_(key_comparator), @@ -126,7 +126,6 @@ Result> KeyValueFileStoreWrite::CreateWriter( std::shared_ptr compact_manager; std::shared_ptr realtime_context_impl; std::optional realtime_store_state; - std::shared_ptr transport_schema; if (realtime_context_) { std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, @@ -134,12 +133,12 @@ Result> KeyValueFileStoreWrite::CreateWriter( partition_map = std::map(partition_values.begin(), partition_values.end()); PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); - transport_schema = RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()); + if (!realtime_schema_layout_) { + return Status::Invalid("PK realtime schema layout is null"); + } auto c_write_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportSchema(*transport_schema, c_write_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema( + *realtime_schema_layout_->StoreWriteSchema(), c_write_schema.get())); PAIMON_ASSIGN_OR_RAISE( RealtimeStoreState store_state, realtime_context_impl->GetOrCreateRealtimeStore( @@ -168,7 +167,7 @@ Result> KeyValueFileStoreWrite::CreateWriter( if (!realtime_context_) { return std::shared_ptr(std::move(writer)); } - return RealtimePrimaryKeyWriter::Create(partition_map, bucket, schema_, transport_schema, + return RealtimePrimaryKeyWriter::Create(partition_map, bucket, realtime_schema_layout_, trimmed_primary_keys, key_comparator_, options_, realtime_context_impl, realtime_store_state.value(), restore_max_seq_number, writer, pool_); @@ -189,9 +188,9 @@ Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { } Status KeyValueFileStoreWrite::Close() { - PAIMON_RETURN_NOT_OK(AbstractFileStoreWrite::Close()); + Status writer_status = AbstractFileStoreWrite::Close(); compact_manager_factory_->Close(); - return Status::OK(); + return writer_status; } } // namespace paimon diff --git a/src/paimon/core/operation/key_value_file_store_write.h b/src/paimon/core/operation/key_value_file_store_write.h index 66c362f2e..c16ab4088 100644 --- a/src/paimon/core/operation/key_value_file_store_write.h +++ b/src/paimon/core/operation/key_value_file_store_write.h @@ -58,6 +58,7 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr& schema_manager, const std::string& commit_user, const std::string& root_path, const std::shared_ptr& table_schema, const std::shared_ptr& schema, + const std::shared_ptr& realtime_schema_layout, const std::shared_ptr& partition_schema, const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, @@ -86,6 +87,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { return realtime_context_ != nullptr; } + std::shared_ptr GetRealtimeContext() const override { + return realtime_context_; + } + private: bool enable_multi_thread_spill_; std::shared_ptr realtime_context_; 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 7977eca36..b71708db3 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 @@ -51,8 +51,7 @@ #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/realtime/realtime_offset_utils.h" -#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/file_store_commit.h" @@ -251,12 +250,12 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { } Result>> - ReadRealtimePrimaryKeyTransportRows( + ReadRealtimePrimaryKeyStoreRows( 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()); + PAIMON_ASSIGN_OR_RAISE(RealtimeReadState read_state, context->AcquireReadState()); + const std::vector& views = read_state.views; if (views.size() != 1) { return Status::Invalid("expected exactly one real-time store"); } @@ -264,12 +263,13 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { 0, arrow::field("id", arrow::int64(), false))), DataField::ConvertDataFieldToArrowField( DataField(1, arrow::field("value", arrow::utf8())))}; - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(arrow::schema(value_fields)); - std::shared_ptr transport_schema = - RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, + arrow::schema(value_fields))); + const std::shared_ptr& store_write_schema = + schema_layout->StoreWriteSchema(); auto c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*transport_schema, c_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*store_write_schema, c_schema.get())); RealtimeQueryContext query_context{c_schema.get(), /*predicate=*/nullptr}; PAIMON_ASSIGN_OR_RAISE( std::vector> readers, @@ -287,12 +287,12 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { std::shared_ptr values = std::dynamic_pointer_cast(array); if (!values || values->num_fields() != 5) { - return Status::Invalid("unexpected realtime primary-key transport batch"); + return Status::Invalid("unexpected real-time primary-key store 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::dynamic_pointer_cast(values->field(0)); + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(values->field(1)); std::shared_ptr offsets = std::dynamic_pointer_cast(values->field(2)); std::shared_ptr ids = @@ -300,7 +300,7 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { std::shared_ptr payloads = std::dynamic_pointer_cast(values->field(4)); if (!row_kinds || !sequences || !offsets || !ids || !payloads) { - return Status::Invalid("unexpected realtime primary-key transport column type"); + return Status::Invalid("unexpected real-time primary-key store column type"); } for (int64_t row = 0; row < values->length(); ++row) { rows.emplace_back(row_kinds->Value(row), ids->Value(row), @@ -432,8 +432,9 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), }); - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, schema)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); std::unique_ptr dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); CreateTable(dir->Str(), schema, options); @@ -467,13 +468,10 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(realtime_schema, R"([[30, 3, "backwards"]])")), "offset moved backwards or was duplicated"); ASSERT_NOK(writer->Write(MakeBatch(realtime_schema, R"([[null, 3, "null-offset"]])"))); - using RealtimePrimaryKeyTransportRow = - std::tuple; - ASSERT_OK_AND_ASSIGN(std::vector transport_rows, - ReadRealtimePrimaryKeyTransportRows(realtime_context)); - ASSERT_EQ((std::vector{ - {0, 1, "old", 0, 10}, {2, 1, "new", 2, 30}, {3, 2, "two", 1, 20}}), - transport_rows); + ASSERT_OK_AND_ASSIGN(auto store_rows, ReadRealtimePrimaryKeyStoreRows(realtime_context)); + ASSERT_EQ( + (decltype(store_rows){{0, 1, "old", 0, 10}, {2, 1, "new", 2, 30}, {3, 2, "two", 1, 20}}), + store_rows); ASSERT_OK_AND_ASSIGN(std::vector progresses, writer->PrepareCommitWithProgress(0)); ASSERT_EQ(1, progresses.size()); @@ -498,8 +496,9 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), }); - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, schema)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); std::unique_ptr dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); CreateTable(dir->Str(), schema, options); @@ -520,13 +519,13 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { const int64_t allocations_before_write = pool->allocation_count; ASSERT_OK(writer->Write(MakeBatch(realtime_schema, R"([[0, 1, "one"]])"))); ASSERT_GT(pool->allocation_count, allocations_before_write); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); ASSERT_OK(writer->Close()); writer.reset(); - using RealtimePrimaryKeyTransportRow = - std::tuple; - ASSERT_OK_AND_ASSIGN(std::vector retained_rows, - ReadRealtimePrimaryKeyTransportRows(realtime_context)); - ASSERT_EQ((std::vector{{0, 1, "one", 0, 0}}), retained_rows); + ASSERT_OK_AND_ASSIGN(auto retained_rows, ReadRealtimePrimaryKeyStoreRows(realtime_context)); + ASSERT_EQ((decltype(retained_rows){{0, 1, "one", 0, 0}}), retained_rows); std::shared_ptr rejecting_pool = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr rejecting_context, @@ -546,8 +545,7 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { ASSERT_NOK_WITH_MSG(rejecting_writer->Write(MakeBatch(realtime_schema, R"([[0, 2, "two"]])")), "Out of memory"); ASSERT_GT(rejecting_pool->allocation_count, rejecting_allocations_before_write); - ASSERT_OK_AND_ASSIGN(std::vector rejected_rows, - ReadRealtimePrimaryKeyTransportRows(rejecting_context)); + ASSERT_OK_AND_ASSIGN(auto rejected_rows, ReadRealtimePrimaryKeyStoreRows(rejecting_context)); ASSERT_TRUE(rejected_rows.empty()); ASSERT_OK(rejecting_writer->Close()); } @@ -560,8 +558,9 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { arrow::field("id", arrow::int64(), false), arrow::field("value", arrow::utf8()), }); - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, schema)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); std::unique_ptr dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); CreateTable(dir->Str(), schema, options); @@ -605,26 +604,23 @@ TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { FileStoreWrite::Create(std::move(write_context))); ASSERT_OK( writer->Write(MakeBatch(realtime_schema, fmt::format(R"([[{}, 1, "legal"]])", max - 1)))); - using RealtimePrimaryKeyTransportRow = - std::tuple; - ASSERT_OK_AND_ASSIGN(std::vector transport_rows, - ReadRealtimePrimaryKeyTransportRows(realtime_context)); - ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), - transport_rows); + ASSERT_OK_AND_ASSIGN(auto store_rows, ReadRealtimePrimaryKeyStoreRows(realtime_context)); + ASSERT_EQ((decltype(store_rows){{0, 1, "legal", max - 1, max - 1}}), store_rows); ASSERT_NOK_WITH_MSG( writer->Write(MakeBatch(realtime_schema, fmt::format(R"([[{}, 2, "overflow"]])", max))), "real-time offset range exceeds INT64_MAX"); - ASSERT_OK_AND_ASSIGN(transport_rows, ReadRealtimePrimaryKeyTransportRows(realtime_context)); - ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), - transport_rows); + ASSERT_OK_AND_ASSIGN(store_rows, ReadRealtimePrimaryKeyStoreRows(realtime_context)); + ASSERT_EQ((decltype(store_rows){{0, 1, "legal", max - 1, max - 1}}), store_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_OK_AND_ASSIGN(RealtimeReadState read_state, context_impl->AcquireReadState()); + ASSERT_EQ(1, read_state.views.size()); ASSERT_EQ(std::optional(OffsetRange(max - 1, max)), - views[0].read_view->GetOffsetRange()); + read_state.views[0].read_view->GetOffsetRange()); + ASSERT_OK_AND_ASSIGN(std::vector boundary_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, boundary_progress.size()); ASSERT_OK(writer->Close()); ASSERT_GE(snapshot_id, 1); } diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index f9175d812..570059abc 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -132,31 +132,42 @@ class MergeFileSplitRead::RealtimeReaderBuilder { const std::vector>& disk_splits, std::vector>&& additional_readers, MergeFileSplitRead* owner) { - ScopeGuard additional_readers_guard([&additional_readers]() { - for (const std::unique_ptr& reader : additional_readers) { - if (reader) { - reader->Close(); - } - } - }); + ScopeGuard additional_readers_guard( + [&additional_readers]() { close_readers_(additional_readers); }); RealtimeReaderBuilder builder(owner); std::vector> readers; + ScopeGuard readers_guard([&readers]() { close_readers_(readers); }); if (!disk_splits.empty()) { - PAIMON_RETURN_NOT_OK(builder.CollectDiskReaders(disk_splits, &readers)); + if (owner->options_.DeletionVectorsEnabled()) { + PAIMON_RETURN_NOT_OK(builder.CollectDiskReaders(disk_splits, &readers)); + } else { + PAIMON_RETURN_NOT_OK(builder.CollectMergedDiskReaders(disk_splits, &readers)); + } } readers.reserve(readers.size() + additional_readers.size()); for (std::unique_ptr& additional_reader : additional_readers) { readers.push_back(std::move(additional_reader)); } additional_readers_guard.Release(); - return builder.CreateMergedReader(std::move(readers)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr result, + builder.CreateMergedReader(std::move(readers))); + readers_guard.Release(); + return result; } private: + inline static const auto close_readers_ = [](const auto& readers) { + for (const auto& reader : readers) { + if (reader) { + reader->Close(); + } + } + }; + explicit RealtimeReaderBuilder(MergeFileSplitRead* owner) : owner_(owner) {} - Status CollectDiskReaders(const std::vector>& disk_splits, - std::vector>* readers) { + Status CollectMergedDiskReaders(const std::vector>& disk_splits, + std::vector>* readers) { std::shared_ptr first_split; std::vector> data_files; std::vector> deletion_files; @@ -189,7 +200,99 @@ class MergeFileSplitRead::RealtimeReaderBuilder { const int32_t bucket = first_split->Bucket(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> merge_function_wrapper, + MergeFileSplitRead::CreateMergeFunctionWrapper(owner_->options_, + owner_->context_->GetTableSchema(), + owner_->value_schema_, owner_->pool_)); + return CollectSectionedDiskReader(data_files, deletion_files, partition, + owner_->predicate_for_keys_, data_file_path_factory, + merge_function_wrapper, readers); + } + + Status CollectDiskReaders(const std::vector>& disk_splits, + std::vector>* readers) { + std::shared_ptr first_split; + std::vector> high_level_files; + std::vector> high_level_deletion_files; + std::vector> level0_files; + std::vector> level0_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) { + return Status::Invalid("merge input disk split is not a data split"); + } + if (!first_split) { + first_split = data_split; + } + 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"); + } + for (size_t i = 0; i < split_files.size(); ++i) { + const std::shared_ptr& file = split_files[i]; + const std::optional& deletion_file = + split_deletion_files.empty() ? std::nullopt : split_deletion_files[i]; + if (file->level == 0) { + level0_files.push_back(file); + level0_deletion_files.push_back(deletion_file); + } else if (file->level > 0) { + high_level_files.push_back(file); + high_level_deletion_files.push_back(deletion_file); + } else { + return Status::Invalid("merge input disk file has a negative level"); + } + } + } + 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)); + + PAIMON_RETURN_NOT_OK(CollectRawDiskReaders(high_level_files, high_level_deletion_files, + partition, owner_->context_->GetPredicate(), + data_file_path_factory, readers)); + return CollectRawDiskReaders(level0_files, level0_deletion_files, partition, + /*predicate=*/nullptr, data_file_path_factory, readers); + } + Status CollectRawDiskReaders(const std::vector>& data_files, + const std::vector>& deletion_files, + const BinaryRow& partition, + const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, + std::vector>* readers) { + if (data_files.empty()) { + return Status::OK(); + } + DeletionVector::Factory dv_factory; + std::vector> disk_sections; + PAIMON_RETURN_NOT_OK( + owner_->CreateDiskSections(data_files, deletion_files, &dv_factory, &disk_sections)); + for (const std::vector& section : disk_sections) { + for (const SortedRun& run : section) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr run_reader, + owner_->CreateReaderForRun(partition, run, dv_factory, predicate, + data_file_path_factory)); + readers->push_back(std::move(run_reader)); + } + } + return Status::OK(); + } + + Status CollectSectionedDiskReader( + const std::vector>& data_files, + const std::vector>& deletion_files, const BinaryRow& partition, + const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, + const std::shared_ptr>& merge_function_wrapper, + std::vector>* readers) { DeletionVector::Factory dv_factory; std::vector> disk_sections; PAIMON_RETURN_NOT_OK( @@ -198,23 +301,14 @@ class MergeFileSplitRead::RealtimeReaderBuilder { return Status::OK(); } std::vector> section_readers; - ScopeGuard section_readers_guard([§ion_readers]() { - for (const std::unique_ptr& reader : section_readers) { - reader->Close(); - } - }); + ScopeGuard section_readers_guard([§ion_readers]() { close_readers_(section_readers); }); section_readers.reserve(disk_sections.size()); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr> merge_function_wrapper, - MergeFileSplitRead::CreateMergeFunctionWrapper(owner_->options_, - owner_->context_->GetTableSchema(), - owner_->value_schema_, owner_->pool_)); for (const std::vector& section : disk_sections) { PAIMON_ASSIGN_OR_RAISE( std::unique_ptr section_reader, owner_->CreateSortMergeReaderForSection( - section, partition, dv_factory, owner_->predicate_for_keys_, - data_file_path_factory, /*drop_delete=*/false, merge_function_wrapper)); + section, partition, dv_factory, predicate, data_file_path_factory, + /*drop_delete=*/false, merge_function_wrapper)); section_readers.push_back( std::make_unique(std::move(section_reader))); } @@ -227,13 +321,7 @@ class MergeFileSplitRead::RealtimeReaderBuilder { Result> CreateMergedReader( std::vector>&& record_readers) { - ScopeGuard record_readers_guard([&record_readers]() { - for (const std::unique_ptr& reader : record_readers) { - if (reader) { - reader->Close(); - } - } - }); + ScopeGuard record_readers_guard([&record_readers]() { close_readers_(record_readers); }); if (record_readers.empty()) { record_readers_guard.Release(); return std::make_unique(std::vector>{}, @@ -247,6 +335,8 @@ class MergeFileSplitRead::RealtimeReaderBuilder { sort_merge_reader->Close(); } }); + // L0 and memory remain unfiltered until MOR is complete. Apply a residual predicate only + // when the read context explicitly enables predicate filtering. PAIMON_ASSIGN_OR_RAISE(std::unique_ptr result, owner_->CreateProjectedReader(std::move(sort_merge_reader), owner_->context_->GetPredicate(), @@ -323,6 +413,16 @@ Result> MergeFileSplitRead::CreateReader( if (!data_split->BeforeFiles().empty()) { return Status::Invalid("this read cannot accept split with before files."); } + const bool contains_level0 = + std::any_of(data_split->DataFiles().begin(), data_split->DataFiles().end(), + [](const std::shared_ptr& file) { return file->level == 0; }); + if (context_->GetRealtimeContext() && options_.DeletionVectorsEnabled() && contains_level0) { + // A disk-only real-time split still needs level-aware predicate pushdown. Route it through + // the same reader used by disk-plus-memory splits so L0 remains an unfiltered durable tail. + std::vector> disk_splits = {split}; + return CreateRealtimeReader(disk_splits, + std::vector>{}); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr data_file_path_factory, path_factory_->CreateDataFilePathFactory(data_split->Partition(), data_split->Bucket())); diff --git a/src/paimon/core/postpone/postpone_bucket_file_store_write.h b/src/paimon/core/postpone/postpone_bucket_file_store_write.h index 1ca65d289..ac21ade87 100644 --- a/src/paimon/core/postpone/postpone_bucket_file_store_write.h +++ b/src/paimon/core/postpone/postpone_bucket_file_store_write.h @@ -110,9 +110,10 @@ class PostponeBucketFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr& executor, const std::shared_ptr& pool) : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, table_schema, schema, - /*write_schema=*/schema, partition_schema, dv_maintainer_factory, - io_manager, options, ignore_previous_files, is_streaming_mode, - ignore_num_bucket_check, executor, pool) {} + /*write_schema=*/schema, + /*realtime_schema_layout=*/nullptr, partition_schema, + dv_maintainer_factory, io_manager, options, ignore_previous_files, + is_streaming_mode, ignore_num_bucket_check, executor, pool) {} Result> CreateWriter( const BinaryRow& partition, int32_t bucket, diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index e4cbf081e..3a819024e 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -30,10 +30,7 @@ #include "paimon/common/data/columnar/columnar_row.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/predicate/predicate_filter.h" -#include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/reader_utils.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/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -66,6 +63,18 @@ bool SupportsMinMax(const std::shared_ptr& type) { } } +Result> ProjectBatch( + const std::shared_ptr& data, + const std::shared_ptr& read_schema, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr projected, + NestedProjectionUtils::AlignArrayToReadType( + data, arrow::struct_(read_schema->fields()), arrow_pool)); + if (!projected || projected->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("memory query projection did not produce a StructArray"); + } + return checked_pointer_cast(projected); +} + } // namespace class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { @@ -74,6 +83,7 @@ class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { : offset_range_(offset_range), batches_(std::move(batches)) { for (const StoredBatch& batch : batches_) { row_count_ += batch.data->length(); + memory_usage_ += batch.memory_usage; } } @@ -90,25 +100,23 @@ class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { } uint64_t GetMemoryUsage() const { - uint64_t result = 0; - for (const StoredBatch& batch : batches_) { - result += batch.memory_usage; - } - return result; + return memory_usage_; } private: OffsetRange offset_range_; std::vector batches_; int64_t row_count_ = 0; + uint64_t memory_usage_ = 0; }; class ArrowRealtimeStore::ReadView : public RealtimeReadView { public: - explicit ReadView(std::vector&& batches) : batches_(std::move(batches)) { - if (!batches_.empty()) { - offset_range_ = - OffsetRange(batches_.front().offset_range.begin, batches_.back().offset_range.end); + explicit ReadView(std::vector>&& segments) + : segments_(std::move(segments)) { + if (!segments_.empty()) { + offset_range_ = OffsetRange(segments_.front()->GetOffsetRange().begin, + segments_.back()->GetOffsetRange().end); } } @@ -116,56 +124,65 @@ class ArrowRealtimeStore::ReadView : public RealtimeReadView { return offset_range_; } - const std::vector& GetBatches() const { - return batches_; + const std::vector>& GetSegments() const { + return segments_; } private: - std::vector batches_; + std::vector> segments_; std::optional offset_range_; }; -class ArrowRealtimeStore::CommitBatchReader : public BatchReader { +class ArrowRealtimeStore::AppendCommitBatchReader : 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()) {} + explicit AppendCommitBatchReader(const std::shared_ptr& segment) + : segment_(segment), metrics_(std::make_shared()) {} Result NextBatch() override { if (!segment_ || next_batch_ >= segment_->GetBatches().size()) { return MakeEofBatch(); } const StoredBatch& stored = segment_->GetBatches()[next_batch_++]; - int64_t row_count = stored.data->length(); + auto c_array = std::make_unique(); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*stored.data, c_array.get(), c_schema.get())); + return ReadBatch(std::move(c_array), std::move(c_schema)); + } - 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) { - PAIMON_RETURN_NOT_OK_FROM_ARROW( - row_kind_builder.Append(static_cast(row_kind))); - } + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + + void Close() override { + segment_.reset(); + } + + private: + std::shared_ptr segment_; + std::shared_ptr metrics_; + size_t next_batch_ = 0; +}; + +class ArrowRealtimeStore::StoredBatchReader : public BatchReader { + public: + StoredBatchReader(const std::shared_ptr& data, + const std::shared_ptr& arrow_pool) + : data_(data), arrow_pool_(arrow_pool), metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!data_) { + return MakeEofBatch(); } - std::shared_ptr row_kind_array; - PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); - - arrow::ArrayVector fields = {row_kind_array}; - fields.insert(fields.end(), stored.data->fields().begin(), stored.data->fields().end()); - arrow::FieldVector schema_fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())}; - const arrow::FieldVector& data_fields = stored.data->struct_type()->fields(); - schema_fields.insert(schema_fields.end(), data_fields.begin(), data_fields.end()); - - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, - arrow::StructArray::Make(fields, schema_fields)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_array, + ArrowUtils::NormalizeArrayOffsets(data_, arrow_pool_.get())); auto c_array = std::make_unique(); auto c_schema = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*result, c_array.get(), c_schema.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*normalized_array, c_array.get(), c_schema.get())); PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); + data_.reset(); + arrow_pool_.reset(); return ReadBatch(std::move(c_array), std::move(c_schema)); } @@ -174,23 +191,24 @@ class ArrowRealtimeStore::CommitBatchReader : public BatchReader { } void Close() override { - segment_.reset(); + data_.reset(); + arrow_pool_.reset(); } private: - std::shared_ptr segment_; + std::shared_ptr data_; std::shared_ptr arrow_pool_; std::shared_ptr metrics_; - size_t next_batch_ = 0; }; -class ArrowRealtimeStore::QueryBatchReader : public BatchReader { +class ArrowRealtimeStore::AppendQueryBatchReader : public BatchReader { public: - QueryBatchReader(const ReadView* view, const std::shared_ptr& read_schema, - const std::shared_ptr& predicate_filter, - std::vector&& statistics_mapping, - const std::shared_ptr& arrow_pool, - const std::shared_ptr& memory_pool) + AppendQueryBatchReader(const std::shared_ptr& view, + const std::shared_ptr& read_schema, + const std::shared_ptr& predicate_filter, + std::vector&& statistics_mapping, + const std::shared_ptr& arrow_pool, + const std::shared_ptr& memory_pool) : view_(view), read_schema_(read_schema), arrow_pool_(arrow_pool), @@ -201,20 +219,28 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { Result NextBatch() override { return Status::Invalid( - "paimon inner reader ArrowRealtimeStore::QueryBatchReader should use " + "paimon inner reader ArrowRealtimeStore::AppendQueryBatchReader should use " "NextBatchWithBitmap"); } Result NextBatchWithBitmap() override { // TODO(xinyu.lxy): Memory query reads return complete stored write batches and // intentionally ignore the configured read batch size. - while (view_ && next_batch_ < view_->GetBatches().size()) { - const StoredBatch& stored = view_->GetBatches()[next_batch_++]; + while (view_ && next_segment_ < view_->GetSegments().size()) { + const std::vector& batches = + view_->GetSegments()[next_segment_]->GetBatches(); + if (next_batch_ >= batches.size()) { + ++next_segment_; + next_batch_ = 0; + continue; + } + const StoredBatch& stored = batches[next_batch_++]; PAIMON_ASSIGN_OR_RAISE(bool may_match, MayMatch(stored)); if (!may_match) { continue; } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, BuildOutput(stored)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, + ProjectBatch(stored.data, read_schema_, arrow_pool_.get())); auto c_array = std::make_unique(); auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( @@ -254,42 +280,31 @@ class ArrowRealtimeStore::QueryBatchReader : public BatchReader { projected_max, projected_null_counts); } - Result> BuildOutput(const StoredBatch& stored) { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr projected, - NestedProjectionUtils::AlignArrayToReadType( - stored.data, arrow::struct_(read_schema_->fields()), arrow_pool_.get())); - if (!projected || projected->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("memory query projection did not produce a StructArray"); - } - std::shared_ptr projected_struct = - checked_pointer_cast(projected); - return projected_struct; - } - private: - const ReadView* view_; + std::shared_ptr view_; std::shared_ptr read_schema_; std::shared_ptr arrow_pool_; std::shared_ptr memory_pool_; std::shared_ptr predicate_filter_; std::vector statistics_mapping_; std::shared_ptr metrics_; + size_t next_segment_ = 0; size_t next_batch_ = 0; }; ArrowRealtimeStore::ArrowRealtimeStore(const std::shared_ptr& write_schema, - StatisticsMode statistics_mode, + RealtimeStoreMode mode, StatisticsMode statistics_mode, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool) : write_schema_(write_schema), memory_pool_(memory_pool), arrow_pool_(arrow_pool), + mode_(mode), statistics_mode_(statistics_mode) {} Result> ArrowRealtimeStore::CollectStatistics( const std::shared_ptr& data) const { - if (statistics_mode_ == StatisticsMode::NONE) { + if (mode_ == RealtimeStoreMode::PRIMARY_KEY || statistics_mode_ == StatisticsMode::NONE) { return std::optional(); } @@ -348,19 +363,13 @@ Result> ArrowRealtimeStore::C } Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { - if (!write_batch.batch) { + if (!write_batch.batch || !write_batch.batch->GetData()) { return Status::Invalid("real-time write batch is null"); } - int64_t row_count = write_batch.batch->GetData()->length; if (write_batch.offset_range.begin < 0 || write_batch.offset_range.begin >= write_batch.offset_range.end) { return Status::Invalid("real-time offset range is invalid"); } - if (!write_batch.batch->GetRowKind().empty() && - static_cast(write_batch.batch->GetRowKind().size()) != row_count) { - return Status::Invalid("real-time row-kind count does not match batch row count"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr data, arrow::ImportArray(write_batch.batch->GetData(), arrow::struct_(write_schema_->fields()))); @@ -383,9 +392,8 @@ Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { ArrowUtils::GetArrayMemoryUsage(statistics->null_counts->data()); } building_memory_usage_ += memory_usage; - building_batches_.push_back( - StoredBatch{std::move(struct_array), write_batch.batch->GetRowKind(), - write_batch.offset_range, std::move(statistics), memory_usage}); + building_batches_.push_back(StoredBatch{std::move(struct_array), write_batch.offset_range, + std::move(statistics), memory_usage}); if (!building_range_) { building_range_ = write_batch.offset_range; } else { @@ -414,19 +422,26 @@ Result>> ArrowRealtimeStore::CreateComm return Status::Invalid("segment was not created by the Arrow real-time store"); } std::vector> readers; - readers.push_back(std::make_unique(arrow_segment, arrow_pool_)); + if (mode_ == RealtimeStoreMode::APPEND_ONLY) { + readers.push_back(std::make_unique(arrow_segment)); + return readers; + } + readers.reserve(arrow_segment->GetBatches().size()); + for (const StoredBatch& batch : arrow_segment->GetBatches()) { + readers.push_back(std::make_unique(batch.data, arrow_pool_)); + } return readers; } Result> ArrowRealtimeStore::AcquireReadView() { std::lock_guard lock(mutex_); - std::vector batches; - for (const std::shared_ptr& segment : sealed_segments_) { - const std::vector& segment_batches = segment->GetBatches(); - batches.insert(batches.end(), segment_batches.begin(), segment_batches.end()); + std::vector> segments = sealed_segments_; + if (!building_batches_.empty()) { + segments.push_back(std::make_shared(building_range_.value(), + std::vector(building_batches_))); } - batches.insert(batches.end(), building_batches_.begin(), building_batches_.end()); - return std::shared_ptr(new ReadView(std::move(batches))); + std::shared_ptr view = std::make_shared(std::move(segments)); + return std::shared_ptr(std::move(view)); } Result>> ArrowRealtimeStore::CreateQueryReaders( @@ -440,6 +455,21 @@ Result>> ArrowRealtimeStore::CreateQuer } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, arrow::ImportSchema(context.read_schema)); + std::vector> readers; + if (!arrow_view->GetOffsetRange()) { + return readers; + } + if (mode_ == RealtimeStoreMode::PRIMARY_KEY) { + for (const std::shared_ptr& segment : arrow_view->GetSegments()) { + for (const StoredBatch& batch : segment->GetBatches()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr projected, + ProjectBatch(batch.data, read_schema, arrow_pool_.get())); + readers.push_back(std::make_unique(projected, arrow_pool_)); + } + } + return readers; + } + std::shared_ptr predicate_filter; if (context.predicate) { predicate_filter = std::dynamic_pointer_cast(context.predicate); @@ -449,14 +479,10 @@ Result>> ArrowRealtimeStore::CreateQuer for (const std::shared_ptr& field : read_schema->fields()) { statistics_mapping.push_back(write_schema_->GetFieldIndex(field->name())); } - std::vector> readers; - if (arrow_view->GetOffsetRange()) { - std::unique_ptr reader = std::make_unique( - arrow_view.get(), read_schema, predicate_filter, std::move(statistics_mapping), - arrow_pool_, memory_pool_); - reader = std::make_unique(std::move(reader), arrow_pool_); - readers.push_back(std::move(reader)); - } + std::unique_ptr reader = std::make_unique( + arrow_view, read_schema, predicate_filter, std::move(statistics_mapping), arrow_pool_, + memory_pool_); + readers.push_back(std::move(reader)); return readers; } diff --git a/src/paimon/core/realtime/arrow_realtime_store.h b/src/paimon/core/realtime/arrow_realtime_store.h index 80cad3cfd..4032568a6 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.h +++ b/src/paimon/core/realtime/arrow_realtime_store.h @@ -38,9 +38,9 @@ namespace paimon { class MemoryPool; /// Internal Arrow-backed implementation of the default `RealtimeStore`. -class ArrowRealtimeStore : public RealtimeStore { +class ArrowRealtimeStore final : public RealtimeStore { public: - ArrowRealtimeStore(const std::shared_ptr& write_schema, + ArrowRealtimeStore(const std::shared_ptr& write_schema, RealtimeStoreMode mode, StatisticsMode statistics_mode, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool); @@ -71,7 +71,6 @@ class ArrowRealtimeStore : public RealtimeStore { struct StoredBatch { std::shared_ptr data; - std::vector row_kinds; OffsetRange offset_range; std::optional statistics; uint64_t memory_usage; @@ -79,8 +78,9 @@ class ArrowRealtimeStore : public RealtimeStore { class Segment; class ReadView; - class CommitBatchReader; - class QueryBatchReader; + class AppendCommitBatchReader; + class StoredBatchReader; + class AppendQueryBatchReader; Result> CollectStatistics( const std::shared_ptr& data) const; @@ -88,6 +88,7 @@ class ArrowRealtimeStore : public RealtimeStore { std::shared_ptr write_schema_; std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; + RealtimeStoreMode mode_; StatisticsMode statistics_mode_; mutable std::mutex mutex_; std::vector building_batches_; diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index dff12b589..d42d37918 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -25,7 +25,6 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/realtime/arrow_realtime_store.h" -#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" namespace paimon { @@ -42,20 +41,14 @@ Result> ArrowRealtimeStoreFactory::Create( } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, arrow::ImportSchema(request.write_schema.get())); - switch (request.mode) { - case RealtimeStoreMode::APPEND_ONLY: { - std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); - return std::make_shared(imported_schema, request.statistics_mode, - request.memory_pool, arrow_pool); - } - case RealtimeStoreMode::PRIMARY_KEY: { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool)); - return std::shared_ptr(std::move(store)); - } + if (request.mode != RealtimeStoreMode::APPEND_ONLY && + request.mode != RealtimeStoreMode::PRIMARY_KEY) { + return Status::Invalid("invalid real-time store mode: ", + static_cast(request.mode)); } - return Status::Invalid("invalid real-time store mode: ", static_cast(request.mode)); + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + return std::make_shared( + imported_schema, request.mode, request.statistics_mode, request.memory_pool, arrow_pool); } } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index 73dc8abb9..70a1a20aa 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -74,7 +74,8 @@ class ArrowRealtimeStoreTest : public testing::Test { } std::shared_ptr CreateStore(StatisticsMode statistics_mode) const { - return std::make_shared(schema_, statistics_mode, pool_, arrow_pool_); + return std::make_shared(schema_, RealtimeStoreMode::APPEND_ONLY, + statistics_mode, pool_, arrow_pool_); } std::unique_ptr MakeBatch(const std::string& json) const { @@ -172,6 +173,19 @@ TEST_F(ArrowRealtimeStoreTest, TestQueryReaderClipsCommittedOffsetWithBitmap) { DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), arrow::field("value", arrow::utf8()), }); + { + std::unique_ptr c_schema = MakeReadSchema(read_schema); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store_->CreateQueryReaders(view, context)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch, + readers[0]->NextBatchWithBitmap()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.first.get(), batch.first.second.get()); + ASSERT_TRUE(import_result.ok()); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(read_schema->fields()))); + } { std::unique_ptr c_schema = MakeReadSchema(read_schema); RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; @@ -235,15 +249,9 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { arrow::ImportArray(batch.first.get(), batch.second.get()); ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); std::shared_ptr actual_array = std::move(import_result).ValueOrDie(); - std::shared_ptr expected_type = arrow::struct_({ - arrow::field("_VALUE_KIND", arrow::int8()), - DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), - arrow::field("id", arrow::int64()), - arrow::field("value", arrow::utf8()), - }); + std::shared_ptr expected_type = arrow::struct_(schema_->fields()); std::shared_ptr expected_array = - arrow::ipc::internal::json::ArrayFromJSON(expected_type, - R"([[0, 0, 1, null], [0, 1, 2, "c"]])") + arrow::ipc::internal::json::ArrayFromJSON(expected_type, R"([[0, 1, null], [1, 2, "c"]])") .ValueOrDie(); ASSERT_TRUE(actual_array->Equals(*expected_array)) << "expected: " << expected_array->ToString() << ", actual: " << actual_array->ToString(); @@ -320,6 +328,7 @@ TEST_F(ArrowRealtimeStoreTest, TestRejectsHandlesFromAnotherStoreImplementation) RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr}; ASSERT_NOK_WITH_MSG(store_->CreateQueryReaders(std::make_shared(), context), "read view was not created by the Arrow real-time store"); + read_schema->release(read_schema.get()); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store_->AcquireReadView()); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp deleted file mode 100644 index f0c485ea9..000000000 --- a/src/paimon/core/realtime/primary_key_realtime_store.cpp +++ /dev/null @@ -1,321 +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_store.h" - -#include -#include -#include -#include -#include - -#include "arrow/api.h" -#include "arrow/c/bridge.h" -#include "arrow/c/helpers.h" -#include "paimon/common/metrics/metrics_impl.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/utils/nested_projection_utils.h" -#include "paimon/macros.h" -#include "paimon/memory/memory_pool.h" - -namespace paimon { - -namespace { - -struct StoredBatch { - std::shared_ptr data; - OffsetRange offset_range; - uint64_t memory_usage; -}; - -class Segment final : public RealtimeSegmentHandle { - public: - Segment(const OffsetRange& range, std::vector&& batches) - : range_(range), batches_(std::move(batches)) { - for (const StoredBatch& batch : batches_) { - row_count_ += batch.data->length(); - } - } - - OffsetRange GetOffsetRange() const override { - return range_; - } - int64_t GetRowCount() const override { - return row_count_; - } - const std::vector& Batches() const { - return batches_; - } - - private: - OffsetRange range_; - std::vector batches_; - int64_t row_count_ = 0; -}; - -class ReadView final : public RealtimeReadView { - public: - 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 range_; - } - - const std::vector>& Segments() const { - return segments_; - } - - private: - std::vector> segments_; - std::optional range_; -}; - -class StoredBatchReader final : public BatchReader { - public: - explicit StoredBatchReader(const StoredBatch& batch, - std::shared_ptr arrow_pool) - : arrow_pool_(std::move(arrow_pool)), - data_(batch.data), - metrics_(std::make_shared()) {} - - Result NextBatch() override { - if (!data_) { - return MakeEofBatch(); - } - auto array = std::make_unique(); - auto schema = std::make_unique(); - ScopeGuard export_guard([array_ptr = array.get(), schema_ptr = schema.get()]() { - ArrowArrayRelease(array_ptr); - ArrowSchemaRelease(schema_ptr); - }); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr record_batch, - arrow::RecordBatch::FromStructArray(data_, arrow_pool_.get())); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr normalized_batch, - ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); - PAIMON_RETURN_NOT_OK_FROM_ARROW( - arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); - PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(array.get(), schema.get(), arrow_pool_)); - data_.reset(); - arrow_pool_.reset(); - export_guard.Release(); - return ReadBatch(std::move(array), std::move(schema)); - } - - std::shared_ptr GetReaderMetrics() const override { - return metrics_; - } - void Close() override { - data_.reset(); - arrow_pool_.reset(); - } - - private: - std::shared_ptr arrow_pool_; - std::shared_ptr data_; - std::shared_ptr metrics_; -}; - -} // namespace - -class PrimaryKeyRealtimeStore::Impl { - public: - Impl(std::shared_ptr transport_schema, - std::shared_ptr arrow_pool) - : transport_schema_(std::move(transport_schema)), arrow_pool_(std::move(arrow_pool)) {} - - 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 (write_batch.offset_range.begin < 0 || - write_batch.offset_range.begin >= write_batch.offset_range.end || row_count <= 0) { - return Status::Invalid("PK real-time offset range is invalid"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr array, - arrow::ImportArray(write_batch.batch->GetData(), - arrow::struct_(transport_schema_->fields()))); - if (!array || array->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("PK real-time transport batch is not a StructArray"); - } - std::shared_ptr transport = - checked_pointer_cast(array); - std::lock_guard lock(mutex_); - if (!building_.empty() && - write_batch.offset_range.begin < building_.back().offset_range.end) { - return Status::Invalid( - "PK real-time offset ranges must be ordered and non-overlapping"); - } - building_.push_back(StoredBatch{transport, write_batch.offset_range, - ArrowUtils::GetArrayMemoryUsage(transport->data())}); - building_memory_usage_ += building_.back().memory_usage; - return Status::OK(); - } - - Result>> SealForCommit() { - std::lock_guard lock(mutex_); - if (building_.empty()) { - return std::optional>(); - } - 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& 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> readers; - readers.reserve(segment->Batches().size()); - for (const StoredBatch& batch : segment->Batches()) { - readers.push_back(std::make_unique(batch, arrow_pool_)); - } - return readers; - } - - Result> AcquireReadView() { - std::lock_guard lock(mutex_); - 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_))); - } - return std::make_shared(std::move(segments)); - } - - Result>> CreateQueryReaders( - const std::shared_ptr& view, 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 == nullptr || context.read_schema->release == nullptr) { - return Status::Invalid("PK real-time query read schema is null"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, - arrow::ImportSchema(context.read_schema)); - std::vector> readers; - for (const std::shared_ptr& segment : typed->Segments()) { - for (const StoredBatch& batch : segment->Batches()) { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr projected, - NestedProjectionUtils::AlignArrayToReadType( - batch.data, arrow::struct_(read_schema->fields()), arrow_pool_.get())); - if (!projected || projected->type_id() != arrow::Type::STRUCT) { - return Status::Invalid( - "PK memory query projection did not produce a StructArray"); - } - StoredBatch query_batch{checked_pointer_cast(projected), - batch.offset_range, /*memory_usage=*/0}; - readers.push_back(std::make_unique(query_batch, arrow_pool_)); - } - } - return readers; - } - - Status AdvanceCommittedOffset(int64_t committed_end_offset) { - std::lock_guard lock(mutex_); - auto first_retained = std::find_if( - sealed_.begin(), sealed_.end(), [committed_end_offset](const auto& segment) { - return segment->GetOffsetRange().end > committed_end_offset; - }); - sealed_.erase(sealed_.begin(), first_retained); - return Status::OK(); - } - - uint64_t GetMemoryUsage() const { - std::lock_guard lock(mutex_); - uint64_t total = building_memory_usage_; - for (const std::shared_ptr& segment : sealed_) { - for (const StoredBatch& batch : segment->Batches()) { - total += batch.memory_usage; - } - } - return total; - } - - private: - std::shared_ptr transport_schema_; - std::shared_ptr arrow_pool_; - mutable std::mutex mutex_; - std::vector building_; - std::vector> sealed_; - uint64_t building_memory_usage_ = 0; -}; - -PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) - : impl_(std::move(impl)) {} -PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; - -Result> PrimaryKeyRealtimeStore::Create( - const std::shared_ptr& transport_schema, - const std::shared_ptr& memory_pool) { - if (!memory_pool) { - return Status::Invalid("PK real-time store memory pool is null"); - } - std::shared_ptr arrow_pool = GetArrowPool(memory_pool); - return std::shared_ptr(new PrimaryKeyRealtimeStore( - std::make_unique(transport_schema, std::move(arrow_pool)))); -} -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, const RealtimeQueryContext& context) { - return impl_->CreateQueryReaders(view, context); -} -Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_end_offset) { - return impl_->AdvanceCommittedOffset(committed_end_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 deleted file mode 100644 index 20a312cf4..000000000 --- a/src/paimon/core/realtime/primary_key_realtime_store.h +++ /dev/null @@ -1,61 +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 - -#include "paimon/realtime/realtime_store.h" - -namespace arrow { -class Schema; -} // namespace arrow - -namespace paimon { - -class MemoryPool; - -/// Internal in-memory implementation of the default primary-key `RealtimeStore`. -class PrimaryKeyRealtimeStore final : public RealtimeStore { - public: - static Result> Create( - const std::shared_ptr& transport_schema, - 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, - const RealtimeQueryContext& context) override; - Status AdvanceCommittedOffset(int64_t committed_end_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 index 10d697ce8..d62f012d2 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -16,8 +16,6 @@ * limitations under the License. */ -#include "paimon/core/realtime/primary_key_realtime_store.h" - #include #include #include @@ -29,10 +27,11 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.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/core/realtime/realtime_offset_utils.h" -#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/macros.h" #include "paimon/memory/memory_pool.h" #include "paimon/realtime/arrow_realtime_store_factory.h" @@ -50,19 +49,20 @@ std::shared_ptr FieldWithId(const std::string& name, ->WithNullable(nullable); } -std::shared_ptr CreateTransportSchema(const arrow::FieldVector& value_fields) { - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(arrow::schema(value_fields)); - return RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()); +std::shared_ptr CreateStoreWriteSchema(const arrow::FieldVector& value_fields) { + EXPECT_OK_AND_ASSIGN( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, arrow::schema(value_fields))); + return schema_layout->StoreWriteSchema(); } -std::shared_ptr TransportSchema() { - return CreateTransportSchema( +std::shared_ptr StoreWriteSchema() { + return CreateStoreWriteSchema( {FieldWithId("id", arrow::int64(), 0), FieldWithId("value", arrow::utf8(), 1)}); } -std::shared_ptr NestedTransportSchema() { - return CreateTransportSchema( +std::shared_ptr NestedStoreWriteSchema() { + return CreateStoreWriteSchema( {FieldWithId("id", arrow::int64(), 0), FieldWithId("value", arrow::struct_({arrow::field("name", arrow::utf8()), @@ -70,10 +70,16 @@ std::shared_ptr NestedTransportSchema() { 1)}); } +std::shared_ptr CreateStore(const std::shared_ptr& schema, + const std::shared_ptr& pool) { + return std::make_shared(schema, RealtimeStoreMode::PRIMARY_KEY, + StatisticsMode::NONE, pool, GetArrowPool(pool)); +} + std::unique_ptr MakeBatch(const std::string& json) { - std::shared_ptr array = - arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(TransportSchema()->fields()), json) - .ValueOrDie(); + std::shared_ptr array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(StoreWriteSchema()->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(); @@ -102,7 +108,7 @@ void AssertOffsetsZero(const ArrowArray* array) { } } -Result ReadJson(std::vector> readers) { +Result> ReadArray(std::vector> readers) { std::vector> batches; for (std::unique_ptr& reader : readers) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, @@ -113,27 +119,35 @@ Result ReadJson(std::vector> readers) } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, arrow::Concatenate(batches)); - return result->ToString(); + return result; +} + +// ReadResultCollector intentionally normalizes struct fields to nullable while decoding +// dictionaries. Use its result type so Array::Equals compares the JSON values, not nullability. +Result> ArrayFromJson(const std::shared_ptr& type, + const std::string& json) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(type, json)); + return array; } TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); 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"([[0, 1, 0, 1, "one"]])"), OffsetRange(0, 0)}), + store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, 0, 0, 1, "one"]])"), OffsetRange(0, 0)}), "offset range is invalid"); ASSERT_OK(store->Write(RealtimeWriteBatch{ - MakeBatch(R"([[0, 1, 0, 1, "one"], [0, 2, 1, 2, "two"]])"), OffsetRange(0, 2)})); + MakeBatch(R"([[1, 0, 0, 1, "one"], [2, 0, 1, 2, "two"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[0, 3, 4, 3, "three"]])"), OffsetRange(4, 5)})); + RealtimeWriteBatch{MakeBatch(R"([[3, 0, 4, 3, "three"]])"), OffsetRange(4, 5)})); ASSERT_NOK_WITH_MSG( - store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)}), + store->Write(RealtimeWriteBatch{MakeBatch(R"([[4, 0, 3, 4, "four"]])"), OffsetRange(3, 4)}), "offset ranges must be ordered and non-overlapping"); ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); @@ -142,29 +156,27 @@ TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { ASSERT_EQ(3, segment.value()->GetRowCount()); ASSERT_GT(store->GetMemoryUsage(), 0); ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[0, 4, 5, 4, "four"]])"), OffsetRange(5, 6)})); + RealtimeWriteBatch{MakeBatch(R"([[4, 0, 5, 4, "four"]])"), OffsetRange(5, 6)})); } TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); ASSERT_OK(store->Write(RealtimeWriteBatch{ - MakeBatch(R"([[1, 6, 1, 1, "before"], [0, 5, 0, 3, "three"]])"), OffsetRange(0, 2)})); + MakeBatch(R"([[6, 1, 1, 1, "before"], [5, 0, 0, 3, "three"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); + RealtimeWriteBatch{MakeBatch(R"([[7, 2, 2, 2, "after"]])"), 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(2, readers.size()); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); - ASSERT_EQ( - "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 0,\n 2\n ]\n-- " - "child 1 type: int64\n [\n 6,\n 5,\n 7\n ]\n-- child 2 type: int64\n [\n " - "1,\n 0,\n 2\n ]\n-- child 3 type: int64\n [\n 1,\n 3,\n 2\n ]\n-- child " - "4 type: string\n [\n \"before\",\n \"three\",\n \"after\"\n ]", - actual); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, ReadArray(std::move(readers))); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + ArrayFromJson(actual->type(), + R"([[6, 1, 1, 1, "before"], [5, 0, 0, 3, "three"], [7, 2, 2, 2, "after"]])")); + ASSERT_TRUE(expected->Equals(actual)); } void AssertSlicedBatch(BatchReader* reader) { @@ -198,13 +210,12 @@ void AssertSlicedBatch(BatchReader* reader) { } TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { - std::shared_ptr schema = NestedTransportSchema(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(schema, GetDefaultPool())); + std::shared_ptr schema = NestedStoreWriteSchema(); + std::shared_ptr store = CreateStore(schema, GetDefaultPool()); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeSlicedBatch( schema, - R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]], [0, 3, 2, 3, ["three", [5, 6]]], [0, 4, 3, 4, ["four", [7, 8]]]])", + R"([[1, 0, 0, 1, ["one", [1, 2]]], [2, 0, 1, 2, ["two", [3, 4]]], [3, 0, 2, 3, ["three", [5, 6]]], [4, 0, 3, 4, ["four", [7, 8]]]])", 1, 2), OffsetRange(0, 2)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, @@ -225,19 +236,18 @@ TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { } TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); ASSERT_OK( 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_OK( - store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 5, 2, "two"]])"), OffsetRange(5, 6)})); + store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, 0, 5, 2, "two"]])"), OffsetRange(5, 6)})); ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); ASSERT_TRUE(segment.has_value()); ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[0, 2, 6, 3, "three"]])"), OffsetRange(6, 7)})); + RealtimeWriteBatch{MakeBatch(R"([[2, 0, 6, 3, "three"]])"), OffsetRange(6, 7)})); ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); ASSERT_TRUE(segment.has_value()); ASSERT_OK_AND_ASSIGN(std::shared_ptr retained_view, store->AcquireReadView()); @@ -262,40 +272,43 @@ TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { ASSERT_FALSE(current_view->GetOffsetRange().has_value()); auto c_schema = std::make_unique(); - ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(retained_view, context)); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); - ASSERT_NE(std::string::npos, actual.find("\"one\"")); - ASSERT_NE(std::string::npos, actual.find("\"two\"")); - ASSERT_NE(std::string::npos, actual.find("\"three\"")); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, ReadArray(std::move(readers))); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + ArrayFromJson(actual->type(), + R"([[0, 0, 4, 1, "one"], [1, 0, 5, 2, "two"], [2, 0, 6, 3, "three"]])")); + ASSERT_TRUE(expected->Equals(actual)); } TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); ASSERT_OK( - store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, 0, 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)})); + store->Write(RealtimeWriteBatch{MakeBatch(R"([[2, 0, 1, 1, "one"]])"), OffsetRange(1, 2)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); auto c_schema = std::make_unique(); - ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); RealtimeQueryContext context{/*read_schema=*/c_schema.get(), /*predicate=*/nullptr}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(2, readers.size()); - ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(std::move(readers))); - ASSERT_NE(std::string::npos, actual.find("\"one\"")); - ASSERT_NE(std::string::npos, actual.find("\"two\"")); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, ReadArray(std::move(readers))); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr expected, + ArrayFromJson(actual->type(), R"([[1, 0, 0, 2, "two"], [2, 0, 1, 1, "one"]])")); + ASSERT_TRUE(expected->Equals(actual)); } TEST(PrimaryKeyRealtimeStoreTest, TestQueryBatchOutlivesStoreAndReader) { - const std::shared_ptr stored_schema = TransportSchema(); + const std::shared_ptr stored_schema = StoreWriteSchema(); std::shared_ptr pool = GetMemoryPool(); std::weak_ptr pool_lifetime = pool; auto write_schema = std::make_unique(); @@ -306,7 +319,7 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryBatchOutlivesStoreAndReader) { std::move(write_schema), /*options=*/{}, pool, RealtimeStoreMode::PRIMARY_KEY})); ASSERT_OK(store->Write( - RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 7, "seven"]])"), OffsetRange(0, 1)})); + RealtimeWriteBatch{MakeBatch(R"([[1, 0, 0, 7, "seven"]])"), OffsetRange(0, 1)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); auto c_schema = std::make_unique(); @@ -346,13 +359,12 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { FieldWithId("profile", arrow::struct_({stored_profile_a}), 1), FieldWithId("items", arrow::list(arrow::struct_({stored_a, stored_b})), 2), FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_x, stored_y})), 3)}; - std::shared_ptr stored_schema = CreateTransportSchema(stored_value_fields); - ASSERT_OK_AND_ASSIGN(std::shared_ptr store, - PrimaryKeyRealtimeStore::Create(stored_schema, GetDefaultPool())); + std::shared_ptr stored_schema = CreateStoreWriteSchema(stored_value_fields); + std::shared_ptr store = CreateStore(stored_schema, GetDefaultPool()); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeSlicedBatch( stored_schema, - R"([[0, 1, 0, 6, [5], [[1, 2]], [["before", [3, 4]]]], [0, 2, 1, 7, [50], [[100, 200], null], [["k1", [7, 8]], ["k2", null]]], [0, 3, 2, 8, [500], [[9, 10]], [["after", [11, 12]]]]])", + R"([[1, 0, 0, 6, [5], [[1, 2]], [["before", [3, 4]]]], [2, 0, 1, 7, [50], [[100, 200], null], [["k1", [7, 8]], ["k2", null]]], [3, 0, 2, 8, [500], [[9, 10]], [["after", [11, 12]]]]])", 1, 1), OffsetRange(0, 1)})); ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); @@ -363,7 +375,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { FieldWithId("items", arrow::list(arrow::struct_({stored_b, stored_a})), 2)); requested_value_fields.push_back( FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_y, stored_x})), 3)); - std::shared_ptr requested_schema = CreateTransportSchema(requested_value_fields); + std::shared_ptr requested_schema = + CreateStoreWriteSchema(requested_value_fields); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*requested_schema, c_schema.get()).ok()); RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 8d45ae8ef..f205d7730 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -37,67 +37,52 @@ #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/append/append_only_writer.h" +#include "paimon/core/core_options.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/realtime/realtime_offset_utils.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/utils/commit_increment.h" #include "paimon/macros.h" #include "paimon/realtime/realtime_context.h" namespace paimon { -namespace { - -Result> AddRealtimeOffsetToSchema(::ArrowSchema* write_schema) { - if (!write_schema || !write_schema->release) { - return Status::Invalid("real-time store write schema is null"); - } - ScopeGuard schema_guard([write_schema]() { ArrowSchemaRelease(write_schema); }); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr input_schema, - arrow::ImportSchema(write_schema)); - schema_guard.Release(); - std::shared_ptr realtime_write_schema = - RealtimeOffsetUtils::CreateInputSchema(input_schema); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*realtime_write_schema, write_schema)); - return realtime_write_schema; -} - -} // namespace Result> RealtimeAppendOnlyWriter::Create( const std::map& partition, int32_t bucket, - 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::map& options, + const std::shared_ptr& schema_layout, const CoreOptions& options, const std::shared_ptr& memory_pool) { if (!realtime_context) { return Status::Invalid("real-time context is null"); } + if (!schema_layout) { + return Status::Invalid("append real-time schema layout is null"); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, RealtimeContextImpl::Cast(realtime_context)); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_write_schema, - AddRealtimeOffsetToSchema(write_schema.get())); - RealtimeStoreCreateRequest request{std::move(write_schema), options, memory_pool, - RealtimeStoreMode::APPEND_ONLY, statistics_mode}; + auto write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*schema_layout->StoreWriteSchema(), write_schema.get())); + RealtimeStoreCreateRequest request{std::move(write_schema), options.ToMap(), memory_pool, + RealtimeStoreMode::APPEND_ONLY, + options.GetRealtimeStoreStatisticsMode()}; PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, realtime_context_impl->GetOrCreateRealtimeStore( std::move(request), RealtimePartitionBucket(partition, bucket))); return std::shared_ptr(new RealtimeAppendOnlyWriter( - store_state.store, file_writer, input_schema, realtime_write_schema, - store_state.initial_offset, memory_pool)); + store_state.store, file_writer, schema_layout, store_state.initial_offset, memory_pool)); } RealtimeAppendOnlyWriter::RealtimeAppendOnlyWriter( const std::shared_ptr& realtime_store, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, - const std::shared_ptr& realtime_write_schema, int64_t next_offset, + const std::shared_ptr& schema_layout, int64_t next_offset, const std::shared_ptr& memory_pool) : arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(realtime_store), file_writer_(file_writer), - input_schema_(input_schema), - realtime_write_schema_(realtime_write_schema), + schema_layout_(schema_layout), next_offset_(next_offset) {} Status RealtimeAppendOnlyWriter::Write(std::unique_ptr&& batch) { @@ -118,31 +103,50 @@ Status RealtimeAppendOnlyWriter::Write(std::unique_ptr&& batch) { return Status::OK(); } std::lock_guard lock(realtime_store_mutex_); - PAIMON_ASSIGN_OR_RAISE( - RealtimeOffsetUtils::ValidatedBatch validated, - RealtimeOffsetUtils::ValidateBatch(batch.get(), realtime_write_schema_, next_offset_)); + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetUtils::ValidatedBatch validated, + RealtimeOffsetUtils::ValidateBatch( + batch.get(), schema_layout_->InputSchema(), next_offset_)); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*validated.data, batch->GetData())); PAIMON_RETURN_NOT_OK( realtime_store_->Write(RealtimeWriteBatch{std::move(batch), validated.offset_range})); next_offset_ = validated.offset_range.end; + has_building_data_ = true; + return Status::OK(); +} + +Status RealtimeAppendOnlyWriter::SealCurrentSegment() { + PAIMON_ASSIGN_OR_RAISE(std::optional> segment, + realtime_store_->SealForCommit()); + if (segment) { + if (!segment.value()) { + return Status::Invalid("append real-time store sealed a null segment"); + } + sealed_segments_.push_back(std::move(segment.value())); + has_building_data_ = false; + } return Status::OK(); } +Status RealtimeAppendOnlyWriter::Seal() { + std::lock_guard lock(realtime_store_mutex_); + return SealCurrentSegment(); +} + Result RealtimeAppendOnlyWriter::PrepareCommit(bool wait_compaction) { std::lock_guard lock(prepare_mutex_); - std::optional> segment; + std::vector> segments; { 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); + PAIMON_RETURN_NOT_OK(SealCurrentSegment()); + segments.swap(sealed_segments_); } - if (segment) { - PAIMON_RETURN_NOT_OK(FlushSegment(segment.value())); + for (const std::shared_ptr& segment : segments) { + PAIMON_RETURN_NOT_OK(FlushSegment(segment)); } PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, file_writer_->PrepareCommit(wait_compaction)); - if (segment) { - increment.SetRealtimeOffsetRange(segment.value()->GetOffsetRange()); + if (!segments.empty()) { + increment.SetRealtimeOffsetRange(OffsetRange(segments.front()->GetOffsetRange().begin, + segments.back()->GetOffsetRange().end)); } return increment; } @@ -168,27 +172,11 @@ Status RealtimeAppendOnlyWriter::FlushSegment( } 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( - "real-time store commit reader must return an INT8 _VALUE_KIND field"); - } - std::shared_ptr row_kinds = - checked_pointer_cast(value_kind); - for (int64_t i = 0; i < row_kinds->length(); ++i) { - if (row_kinds->IsNull(i) || - row_kinds->Value(i) != static_cast(RecordBatch::RowKind::INSERT)) { - return Status::Invalid( - "append real-time store commit reader returned a non-INSERT row"); - } - } - PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( - struct_array, SpecialFields::ValueKind().Name())); PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( struct_array, SpecialFields::RealtimeOffset().Name())); - if (!struct_array->type()->Equals(arrow::struct_(input_schema_->fields()))) { + if (!struct_array->type()->Equals( + arrow::struct_(schema_layout_->CommitSchema()->fields()))) { return Status::Invalid( "real-time store commit reader schema does not match table write schema"); } @@ -240,6 +228,11 @@ Status RealtimeAppendOnlyWriter::Close() { return file_writer_->Close(); } +bool RealtimeAppendOnlyWriter::HasUnpreparedRealtimeData() const { + std::lock_guard lock(realtime_store_mutex_); + return has_building_data_ || !sealed_segments_.empty(); +} + std::shared_ptr RealtimeAppendOnlyWriter::GetMetrics() const { return file_writer_->GetMetrics(); } diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h index 22ebed560..cf27245f1 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.h +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -24,12 +24,12 @@ #include #include #include +#include +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/utils/batch_writer.h" #include "paimon/realtime/realtime_store.h" -struct ArrowSchema; - namespace arrow { class MemoryPool; class Schema; @@ -38,6 +38,7 @@ class Schema; namespace paimon { class AppendOnlyWriter; +class CoreOptions; class MemoryPool; class RealtimeContext; @@ -45,15 +46,15 @@ class RealtimeAppendOnlyWriter : public BatchWriter { public: static Result> Create( const std::map& partition, int32_t bucket, - 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::map& options, + const std::shared_ptr& schema_layout, const CoreOptions& options, const std::shared_ptr& memory_pool); Status Write(std::unique_ptr&& batch) override; + Status Seal() override; + Result PrepareCommit(bool wait_compaction) override; Status Compact(bool full_compaction) override; @@ -68,24 +69,27 @@ class RealtimeAppendOnlyWriter : public BatchWriter { Status Close() override; + bool HasUnpreparedRealtimeData() const override; + std::shared_ptr GetMetrics() const override; private: RealtimeAppendOnlyWriter(const std::shared_ptr& realtime_store, const std::shared_ptr& file_writer, - const std::shared_ptr& input_schema, - const std::shared_ptr& realtime_write_schema, + const std::shared_ptr& schema_layout, int64_t next_offset, const std::shared_ptr& memory_pool); Status FlushSegment(const std::shared_ptr& segment); + Status SealCurrentSegment(); std::shared_ptr arrow_pool_; std::shared_ptr realtime_store_; std::shared_ptr file_writer_; - std::shared_ptr input_schema_; - std::shared_ptr realtime_write_schema_; + std::shared_ptr schema_layout_; + std::vector> sealed_segments_; int64_t next_offset_; - std::mutex realtime_store_mutex_; + bool has_building_data_ = false; + mutable std::mutex realtime_store_mutex_; std::mutex prepare_mutex_; }; diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index 2b9d7d07c..0f1a4cf8b 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -96,8 +96,23 @@ Status RealtimeContextImpl::Start() { return Status::OK(); } +void RealtimeContextImpl::Invalidate() { + invalidated_.store(true); +} + +Status RealtimeContextImpl::CheckUsable() const { + if (invalidated_.load()) { + return Status::Invalid( + "real-time context cannot be reused after a writer failure or an unprepared writer " + "close; create a new RealtimeContext and FileStoreWrite, then let upstream recover " + "input from the durable recovery offset persisted in the snapshot"); + } + return Status::OK(); +} + Result RealtimeContextImpl::GetOrCreateRealtimeStore( RealtimeStoreCreateRequest&& request, const RealtimePartitionBucket& partition_bucket) { + PAIMON_RETURN_NOT_OK(CheckUsable()); if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("real-time store write schema is null"); } @@ -163,6 +178,7 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( Result RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + PAIMON_RETURN_NOT_OK(CheckUsable()); std::lock_guard lock(mutex_); auto iter = stores_.find(partition_bucket); if (iter == stores_.end()) { @@ -177,17 +193,20 @@ Result RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( return entry.materialized_max_sequence_number; } -Result> RealtimeContextImpl::AcquireReadViews() { - std::lock_guard lock(mutex_); - std::vector result; - result.reserve(stores_.size()); +Result RealtimeContextImpl::AcquireReadState() { + PAIMON_RETURN_NOT_OK(CheckUsable()); + std::lock_guard progress_lock(progress_mutex_); + std::lock_guard registry_lock(mutex_); + RealtimeReadState result; + result.views.reserve(stores_.size()); + result.committed_offsets = committed_offsets_; for (const auto& [partition_bucket, store] : stores_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, store.store->AcquireReadView()); if (!read_view) { return Status::Invalid("real-time store returned a null read view"); } - result.push_back( + result.views.push_back( RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; @@ -195,6 +214,7 @@ Result> RealtimeContextImpl::AcquireRea Result RealtimeContextImpl::PinReadView(const RealtimePartitionBucketView& view, int64_t ttl_millis) { + PAIMON_RETURN_NOT_OK(CheckUsable()); if (!view.store || !view.read_view) { return Status::Invalid("cannot pin an incomplete real-time read view"); } @@ -269,6 +289,7 @@ Status RealtimeContextImpl::ReleaseReadView(const std::string& opaque_ticket) { Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, const RealtimeOffsetMap& committed_offsets) { + PAIMON_RETURN_NOT_OK(CheckUsable()); if (snapshot_id < 0) { return Status::Invalid("real-time refresh snapshot id must not be negative"); } diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index ea069a5cd..c8f8bbaa0 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include #include @@ -59,6 +60,11 @@ struct RealtimePartitionBucketView { std::shared_ptr read_view; }; +struct RealtimeReadState { + std::vector views; + RealtimeOffsetMap committed_offsets; +}; + class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { public: static Result> Create( @@ -69,13 +75,18 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { static Result> Cast( const std::shared_ptr& context); + // Fences new writers and scans while allowing already pinned views to finish. + void Invalidate(); + + Status CheckUsable() const; + Result GetOrCreateRealtimeStore( RealtimeStoreCreateRequest&& request, const RealtimePartitionBucket& partition_bucket); Result AdvanceMaterializedMaxSequenceNumber( const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number); - Result> AcquireReadViews(); + Result AcquireReadState(); Result PinReadView(const RealtimePartitionBucketView& view, int64_t ttl_millis); @@ -123,6 +134,7 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::condition_variable read_views_cv_; std::map pinned_read_views_; std::deque read_view_release_queue_; + std::atomic invalidated_{false}; bool stopping_ = false; std::thread read_view_cleanup_thread_; }; diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 78f590e07..405dc1008 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -159,8 +159,8 @@ TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { ASSERT_NE(first.store, fourth.store); ASSERT_EQ(3, factory->stores.size()); - ASSERT_OK_AND_ASSIGN(std::vector views, - context->AcquireReadViews()); + ASSERT_OK_AND_ASSIGN(RealtimeReadState read_state, context->AcquireReadState()); + const std::vector& views = read_state.views; ASSERT_EQ(3, views.size()); const RealtimePartitionBucket expected_partition_bucket({{"dt", "2026-08-02"}}, 0); ASSERT_EQ(expected_partition_bucket, views[0].partition_bucket); @@ -376,8 +376,8 @@ TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, GetDefaultPool())); - ASSERT_OK_AND_ASSIGN(std::vector views, - context->AcquireReadViews()); + ASSERT_OK_AND_ASSIGN(RealtimeReadState read_state, context->AcquireReadState()); + const std::vector& views = read_state.views; ASSERT_EQ(1, views.size()); ASSERT_NOK_WITH_MSG(context->PinReadView(views[0], /*ttl_millis=*/0), @@ -400,8 +400,8 @@ TEST(RealtimeContextTest, TestExpiresAbandonedReadViewTicket) { ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, GetDefaultPool())); - ASSERT_OK_AND_ASSIGN(std::vector views, - context->AcquireReadViews()); + ASSERT_OK_AND_ASSIGN(RealtimeReadState read_state, context->AcquireReadState()); + std::vector& views = read_state.views; ASSERT_EQ(1, views.size()); std::weak_ptr weak_view = views[0].read_view; ASSERT_OK_AND_ASSIGN(std::string ticket, context->PinReadView(views[0], /*ttl_millis=*/10)); @@ -432,7 +432,7 @@ TEST(RealtimeContextTest, TestRejectsNullPluginResults) { ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, GetDefaultPool())); factory->stores[0]->return_null_read_view = true; - ASSERT_NOK_WITH_MSG(context->AcquireReadViews(), "real-time store returned a null read view"); + ASSERT_NOK_WITH_MSG(context->AcquireReadState(), "real-time store returned a null read view"); } } // namespace diff --git a/src/paimon/core/realtime/realtime_offset_utils.h b/src/paimon/core/realtime/realtime_offset_utils.h index 23bf429df..232044a3b 100644 --- a/src/paimon/core/realtime/realtime_offset_utils.h +++ b/src/paimon/core/realtime/realtime_offset_utils.h @@ -46,14 +46,6 @@ class RealtimeOffsetUtils { OffsetRange offset_range; }; - static std::shared_ptr CreateInputSchema( - const std::shared_ptr& write_schema) { - arrow::FieldVector fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; - fields.insert(fields.end(), write_schema->fields().begin(), write_schema->fields().end()); - return arrow::schema(std::move(fields), write_schema->metadata()); - } - static Result ValidateBatch( RecordBatch* batch, const std::shared_ptr& realtime_input_schema, int64_t minimum_offset) { diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.cpp b/src/paimon/core/realtime/realtime_primary_key_reader.cpp index d21977764..e3c78b331 100644 --- a/src/paimon/core/realtime/realtime_primary_key_reader.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_reader.cpp @@ -22,9 +22,6 @@ #include #include -#include "arrow/type.h" -#include "paimon/common/table/special_fields.h" -#include "paimon/common/types/data_field.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/key_value_data_file_record_reader.h" #include "paimon/core/key_value.h" @@ -68,16 +65,6 @@ Result>> CreateKeyValueReaders } // namespace -std::shared_ptr RealtimePrimaryKeyLayout::CreateSchema( - const std::vector>& value_fields) { - arrow::FieldVector fields = { - DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), - DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) - ->WithNullable(false)}; - fields.insert(fields.end(), value_fields.begin(), value_fields.end()); - return arrow::schema(std::move(fields)); -} - Result>> RealtimePrimaryKeyReaderFactory::CreateForCommit( std::vector>&& readers, diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.h b/src/paimon/core/realtime/realtime_primary_key_reader.h index d5b039392..837278907 100644 --- a/src/paimon/core/realtime/realtime_primary_key_reader.h +++ b/src/paimon/core/realtime/realtime_primary_key_reader.h @@ -31,17 +31,6 @@ class BatchReader; class MemoryPool; class RealtimeStoreReadPipeline; -/// Creates the Arrow schema used for PK realtime transport batches. -class RealtimePrimaryKeyLayout { - public: - RealtimePrimaryKeyLayout() = delete; - ~RealtimePrimaryKeyLayout() = delete; - - /// Creates `_VALUE_KIND`, `_SEQUENCE_NUMBER`, then value fields. - static std::shared_ptr CreateSchema( - const std::vector>& value_fields); -}; - class RealtimePrimaryKeyReaderFactory { public: RealtimePrimaryKeyReaderFactory() = delete; diff --git a/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp index 1b9dcb6c9..d365de504 100644 --- a/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp @@ -31,7 +31,7 @@ #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" -#include "paimon/core/realtime/realtime_offset_utils.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/realtime/realtime_store_read_pipeline.h" #include "paimon/memory/memory_pool.h" #include "paimon/realtime/offset_range.h" @@ -52,10 +52,11 @@ std::shared_ptr MakeField(const std::string& name, DataField(field_id, arrow::field(name, type, nullable))); } -std::shared_ptr MakeTransportSchema(const arrow::FieldVector& value_fields) { - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(arrow::schema(value_fields)); - return RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()); +std::shared_ptr MakeStoreWriteSchema(const arrow::FieldVector& value_fields) { + EXPECT_OK_AND_ASSIGN( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, arrow::schema(value_fields))); + return schema_layout->StoreWriteSchema(); } Result>> @@ -64,12 +65,12 @@ CreateRealtimePrimaryKeyQueryReadersForTest(std::vector& key_schema, const std::shared_ptr& value_schema, const std::shared_ptr& memory_pool) { - std::shared_ptr write_schema = MakeTransportSchema(value_schema->fields()); - std::shared_ptr logical_schema = - RealtimePrimaryKeyLayout::CreateSchema(value_schema->fields()); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, value_schema)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr pipeline, - RealtimeStoreReadPipeline::Create(logical_schema, write_schema, memory_pool, + RealtimeStoreReadPipeline::Create(schema_layout->QuerySchema(), *schema_layout, memory_pool, GetArrowPool(memory_pool))); return RealtimePrimaryKeyReaderFactory::CreateForQuery( std::move(readers), visible_offsets, key_schema, value_schema, memory_pool, *pipeline); @@ -138,50 +139,28 @@ class RealtimePrimaryKeyReaderTest : public testing::Test { std::shared_ptr pool_ = GetDefaultPool(); }; -TEST_F(RealtimePrimaryKeyReaderTest, TestPrimaryKeySchemaLayouts) { - arrow::FieldVector value_fields = {arrow::field("key", arrow::int64(), false), - arrow::field("value", arrow::utf8())}; - std::shared_ptr schema = MakeTransportSchema(value_fields); - - ASSERT_EQ(schema->field(0)->name(), "_VALUE_KIND"); - ASSERT_EQ(schema->field(1)->name(), "_SEQUENCE_NUMBER"); - ASSERT_EQ(schema->field(2)->name(), "_REALTIME_OFFSET"); - ASSERT_EQ(schema->field(3)->name(), "key"); - ASSERT_EQ(schema->field(4)->name(), "value"); - ASSERT_FALSE(schema->field(0)->nullable()); - ASSERT_FALSE(schema->field(1)->nullable()); - ASSERT_EQ(schema->field(2)->nullable(), SpecialFields::RealtimeOffset().Nullable()); - ASSERT_FALSE(schema->field(3)->nullable()); - ASSERT_TRUE(schema->field(4)->nullable()); - - std::shared_ptr logical_schema = - RealtimePrimaryKeyLayout::CreateSchema(value_fields); - ASSERT_EQ(logical_schema->field(0)->name(), "_VALUE_KIND"); - ASSERT_EQ(logical_schema->field(1)->name(), "_SEQUENCE_NUMBER"); - ASSERT_EQ(logical_schema->field(2)->name(), "key"); - ASSERT_EQ(logical_schema->field(3)->name(), "value"); -} - TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsCommittedPrefix) { 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 transport_schema = MakeTransportSchema(value_schema->fields()); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); - auto transport_array = std::dynamic_pointer_cast( - arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ - [0, 100, 0, 1, 10], - [0, 101, 1, 2, 20], - [0, 102, 2, 4, 40], - [0, 103, 3, 6, 60] + std::shared_ptr store_write_schema = + MakeStoreWriteSchema(value_schema->fields()); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); + auto store_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, R"([ + [100, 0, 0, 1, 10], + [101, 0, 1, 2, 20], + [102, 0, 2, 4, 40], + [103, 0, 3, 6, 60] ])") .ValueOrDie()); std::vector> batch_readers; batch_readers.push_back( - std::make_unique(transport_array, transport_type, 2)); + std::make_unique(store_array, store_write_type, 2)); ASSERT_OK_AND_ASSIGN( std::vector> readers, CreateRealtimePrimaryKeyQueryReadersForTest(std::move(batch_readers), OffsetRange(2, 4), @@ -203,21 +182,22 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsCommittedPrefix) { TEST_F(RealtimePrimaryKeyReaderTest, TestQueryOffsetCoverageAcrossReadersAndBatches) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); - std::shared_ptr transport_schema = MakeTransportSchema({key}); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr store_write_schema = MakeStoreWriteSchema({key}); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); std::shared_ptr first_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, - R"([[0, 10, 2, 1], [0, 11, 0, 2]])") + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, + R"([[10, 0, 2, 1], [11, 0, 0, 2]])") .ValueOrDie(); std::shared_ptr second_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, - R"([[0, 12, 3, 3], [0, 13, 1, 4]])") + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, + R"([[12, 0, 3, 3], [13, 0, 1, 4]])") .ValueOrDie(); std::vector> batch_readers; - batch_readers.push_back( - std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); - batch_readers.push_back( - std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back(std::make_unique(first_array, store_write_type, + /*read_batch_size=*/1)); + batch_readers.push_back(std::make_unique(second_array, store_write_type, + /*read_batch_size=*/1)); ASSERT_OK_AND_ASSIGN( std::vector> readers, @@ -249,13 +229,14 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsEmptyReadersForEmptyVisibleR TEST_F(RealtimePrimaryKeyReaderTest, TestQueryBitmapBounds) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); - std::shared_ptr transport_schema = MakeTransportSchema({key}); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); - std::shared_ptr transport_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1]])") + std::shared_ptr store_write_schema = MakeStoreWriteSchema({key}); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); + std::shared_ptr store_array = + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, R"([[10, 0, 0, 1]])") .ValueOrDie(); auto batch_reader = std::make_unique( - std::make_unique(transport_array, transport_type, /*batch_size=*/1), + std::make_unique(store_array, store_write_type, /*batch_size=*/1), /*row_id=*/1); ASSERT_OK_AND_ASSIGN( @@ -273,16 +254,17 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestQueryBitmapBounds) { TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsPartialBitmap) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); - std::shared_ptr transport_schema = MakeTransportSchema({key}); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); - std::shared_ptr transport_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, - R"([[0, 10, 0, 1], [0, 11, 1, 2]])") + std::shared_ptr store_write_schema = MakeStoreWriteSchema({key}); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); + std::shared_ptr store_array = + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, + R"([[10, 0, 0, 1], [11, 0, 1, 2]])") .ValueOrDie(); RoaringBitmap32 partial_bitmap; partial_bitmap.Add(0); auto batch_reader = std::make_unique( - transport_array, transport_type, partial_bitmap, /*read_batch_size=*/2); + store_array, store_write_type, partial_bitmap, /*read_batch_size=*/2); batch_reader->EnableRandomizeBatchSize(false); ASSERT_OK_AND_ASSIGN( @@ -295,21 +277,22 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsPartialBitmap) { TEST_F(RealtimePrimaryKeyReaderTest, TestCommitCoverageAcrossReadersAndBatches) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); - std::shared_ptr transport_schema = MakeTransportSchema({key}); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr store_write_schema = MakeStoreWriteSchema({key}); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); std::shared_ptr first_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, - R"([[0, 10, 2, 1], [0, 11, 0, 3]])") + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, + R"([[10, 0, 2, 1], [11, 0, 0, 3]])") .ValueOrDie(); std::shared_ptr second_array = - arrow::ipc::internal::json::ArrayFromJSON(transport_type, - R"([[0, 12, 1, 2], [0, 13, 3, 4]])") + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, + R"([[12, 0, 1, 2], [13, 0, 3, 4]])") .ValueOrDie(); std::vector> batch_readers; - batch_readers.push_back( - std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); - batch_readers.push_back( - std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back(std::make_unique(first_array, store_write_type, + /*read_batch_size=*/1)); + batch_readers.push_back(std::make_unique(second_array, store_write_type, + /*read_batch_size=*/1)); ASSERT_OK_AND_ASSIGN(std::vector> readers, RealtimePrimaryKeyReaderFactory::CreateForCommit( @@ -329,10 +312,10 @@ TEST_F(RealtimePrimaryKeyReaderTest, 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 actual_schema = MakeTransportSchema({key}); + std::shared_ptr actual_schema = MakeStoreWriteSchema({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(); + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[10, 0, 0, 1]])").ValueOrDie(); auto batch_reader = std::make_unique(actual, actual_type, 1); ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, @@ -344,14 +327,14 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestBadCommitBatch) { TEST_F(RealtimePrimaryKeyReaderTest, TestSafeDecode) { std::shared_ptr key = MakeField("key", arrow::int32(), 0); std::shared_ptr value_schema = arrow::schema({key}); - std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr store_write_schema = MakeStoreWriteSchema({key}); - arrow::FieldVector invalid_fields = transport_schema->fields(); - invalid_fields[0] = invalid_fields[0]->WithName("wrong_value_kind"); + arrow::FieldVector invalid_fields = store_write_schema->fields(); + invalid_fields[1] = invalid_fields[1]->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()); + arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[10, 0, 0, 1]])").ValueOrDie()); auto batch_reader = std::make_unique(invalid_array, invalid_type, 1); ASSERT_OK_AND_ASSIGN( @@ -384,16 +367,17 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestNestedValues) { 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}); - std::shared_ptr transport_schema = - MakeTransportSchema(query_value_schema->fields()); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); - std::shared_ptr transport_array = + std::shared_ptr store_write_schema = + MakeStoreWriteSchema(query_value_schema->fields()); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); + std::shared_ptr store_array = arrow::ipc::internal::json::ArrayFromJSON( - transport_type, - R"([[0, 10, 0, 1, [[200, 100], [400, 300]], [["k1", [8, 7]], ["k2", [10, 9]]], [[[12, 11], 13], [[22, 21], 23]]]])") + store_write_type, + R"([[10, 0, 0, 1, [[200, 100], [400, 300]], [["k1", [8, 7]], ["k2", [10, 9]]], [[[12, 11], 13], [[22, 21], 23]]]])") .ValueOrDie(); - auto batch_reader = std::make_unique(transport_array, transport_type, 1); + auto batch_reader = std::make_unique(store_array, store_write_type, 1); ASSERT_OK_AND_ASSIGN( std::unique_ptr reader, CreateRealtimePrimaryKeyQueryReaderForTest(std::move(batch_reader), OffsetRange(0, 1), @@ -449,17 +433,19 @@ TEST_F(RealtimePrimaryKeyReaderTest, TestFactoryRejectsNullReader) { std::shared_ptr value_schema = DataField::ConvertDataFieldsToArrowSchema(value_fields); std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); - std::shared_ptr transport_schema = MakeTransportSchema(value_schema->fields()); - std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); - auto transport_array = std::dynamic_pointer_cast( - arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ - [0, 10, 0, 1, 100] + std::shared_ptr store_write_schema = + MakeStoreWriteSchema(value_schema->fields()); + std::shared_ptr store_write_type = + arrow::struct_(store_write_schema->fields()); + auto store_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(store_write_type, R"([ + [10, 0, 0, 1, 100] ])") .ValueOrDie()); std::vector> batch_readers; batch_readers.push_back( - std::make_unique(transport_array, transport_type, 1)); + std::make_unique(store_array, store_write_type, 1)); batch_readers.push_back(nullptr); ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyReaderFactory::CreateForCommit( std::move(batch_readers), key_schema, value_schema, pool_), diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index c3367b62a..1afe8781b 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -48,10 +48,10 @@ namespace paimon { namespace { -Result CreateRealtimePrimaryKeyTransportBatch( +Result CreateRealtimePrimaryKeyStoreBatch( RealtimeOffsetUtils::ValidatedBatch&& validated, const std::vector& row_kinds, - const std::shared_ptr& transport_schema, + const std::shared_ptr& store_write_schema, const std::vector& trimmed_primary_keys, int64_t first_sequence_number, arrow::MemoryPool* arrow_pool) { std::shared_ptr values = std::move(validated.data); @@ -70,11 +70,11 @@ Result CreateRealtimePrimaryKeyTransportBat std::shared_ptr sequence_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); - arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array)}; + arrow::ArrayVector columns = {std::move(sequence_array), std::move(kind_array)}; columns.insert(columns.end(), values->fields().begin(), values->fields().end()); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr transport, - arrow::StructArray::Make(std::move(columns), transport_schema->fields())); + std::shared_ptr store_batch, + arrow::StructArray::Make(std::move(columns), store_write_schema->fields())); std::vector sort_keys; sort_keys.reserve(trimmed_primary_keys.size() + 1); @@ -87,10 +87,10 @@ Result CreateRealtimePrimaryKeyTransportBat arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( arrow::Datum indices, - arrow::compute::SortIndices(arrow::Datum(transport), options, &context)); + arrow::compute::SortIndices(arrow::Datum(store_batch), options, &context)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( arrow::Datum sorted, - arrow::compute::Take(arrow::Datum(transport), indices, + arrow::compute::Take(arrow::Datum(store_batch), indices, arrow::compute::TakeOptions::NoBoundsCheck(), &context)); validated.data = checked_pointer_cast(sorted.make_array()); return validated; @@ -100,8 +100,7 @@ Result CreateRealtimePrimaryKeyTransportBat Result> RealtimePrimaryKeyWriter::Create( const std::map& partition, int32_t bucket, - const std::shared_ptr& write_schema, - const std::shared_ptr& transport_schema, + const std::shared_ptr& schema_layout, const std::vector& trimmed_primary_keys, const std::shared_ptr& key_comparator, const CoreOptions& options, const std::shared_ptr& realtime_context, @@ -115,10 +114,13 @@ Result> RealtimePrimaryKeyWriter::Crea if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); } + if (!schema_layout) { + return Status::Invalid("PK realtime schema layout is null"); + } 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); + std::shared_ptr field = schema_layout->UserSchema()->GetFieldByName(key); if (!field) { return Status::Invalid("PK field is missing from write schema: ", key); } @@ -129,10 +131,9 @@ Result> RealtimePrimaryKeyWriter::Crea realtime_context->AdvanceMaterializedMaxSequenceNumber( partition_bucket, restored_max_sequence_number)); return std::shared_ptr(new RealtimePrimaryKeyWriter( - store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, - transport_schema, arrow::schema(std::move(key_fields)), trimmed_primary_keys, - key_comparator, options, store_state.initial_offset, initial_max_sequence_number, - memory_pool)); + store_state.store, merge_tree_writer, realtime_context, partition_bucket, schema_layout, + arrow::schema(std::move(key_fields)), trimmed_primary_keys, key_comparator, options, + store_state.initial_offset, initial_max_sequence_number, memory_pool)); } RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( @@ -140,8 +141,7 @@ RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( 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& transport_schema, + const std::shared_ptr& schema_layout, const std::shared_ptr& key_schema, const std::vector& trimmed_primary_keys, const std::shared_ptr& key_comparator, const CoreOptions& options, @@ -153,9 +153,7 @@ RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( merge_tree_writer_(merge_tree_writer), realtime_context_(realtime_context), partition_bucket_(partition_bucket), - write_schema_(write_schema), - realtime_input_schema_(RealtimeOffsetUtils::CreateInputSchema(write_schema)), - transport_schema_(transport_schema), + schema_layout_(schema_layout), key_schema_(key_schema), trimmed_primary_keys_(trimmed_primary_keys), key_comparator_(key_comparator), @@ -181,56 +179,67 @@ Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { static_cast(validated); } std::lock_guard lock(realtime_store_mutex_); - PAIMON_ASSIGN_OR_RAISE( - RealtimeOffsetUtils::ValidatedBatch validated, - RealtimeOffsetUtils::ValidateBatch(batch.get(), realtime_input_schema_, next_offset_)); + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetUtils::ValidatedBatch validated, + RealtimeOffsetUtils::ValidateBatch( + batch.get(), schema_layout_->InputSchema(), next_offset_)); // 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(RealtimeOffsetUtils::ValidatedBatch transport, - CreateRealtimePrimaryKeyTransportBatch( - std::move(validated), row_kinds, transport_schema_, + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetUtils::ValidatedBatch store_batch, + CreateRealtimePrimaryKeyStoreBatch( + std::move(validated), row_kinds, schema_layout_->StoreWriteSchema(), trimmed_primary_keys_, first_sequence, arrow_pool_.get())); auto output = std::make_unique(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport.data, output.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*store_batch.data, output.get())); PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(output.get(), /*schema=*/nullptr, arrow_pool_)); RecordBatchBuilder builder(output.get()); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr store_record_batch, builder.Finish()); PAIMON_RETURN_NOT_OK(realtime_store_->Write( - RealtimeWriteBatch{std::move(transport_batch), transport.offset_range})); - next_offset_ = transport.offset_range.end; + RealtimeWriteBatch{std::move(store_record_batch), store_batch.offset_range})); + next_offset_ = store_batch.offset_range.end; last_sequence_number_ += count; + has_building_data_ = true; PAIMON_RETURN_NOT_OK(realtime_context_->AdvanceMaterializedMaxSequenceNumber( partition_bucket_, last_sequence_number_)); return Status::OK(); } +Status RealtimePrimaryKeyWriter::SealCurrentSegment() { + PAIMON_ASSIGN_OR_RAISE(std::optional> segment, + realtime_store_->SealForCommit()); + if (segment) { + if (!segment.value()) { + return Status::Invalid("PK real-time store sealed a null segment"); + } + sealed_segments_.push_back(std::move(segment.value())); + has_building_data_ = false; + } + return Status::OK(); +} + +Status RealtimePrimaryKeyWriter::Seal() { + std::lock_guard lock(realtime_store_mutex_); + return SealCurrentSegment(); +} + Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { std::lock_guard prepare_lock(prepare_mutex_); - std::optional> segment; + std::vector> segments; { std::lock_guard store_lock(realtime_store_mutex_); - PAIMON_ASSIGN_OR_RAISE(std::optional> sealed, - realtime_store_->SealForCommit()); - segment = std::move(sealed); + PAIMON_RETURN_NOT_OK(SealCurrentSegment()); + segments.swap(sealed_segments_); } - if (segment && !segment.value()) { - return Status::Invalid("PK real-time store sealed a null segment"); - } - std::optional sealed_range; - if (segment) { - sealed_range = segment.value()->GetOffsetRange(); - 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())); + for (const std::shared_ptr& segment : segments) { + PAIMON_RETURN_NOT_OK(FlushSegment(segment)); } PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, merge_tree_writer_->PrepareCommit(wait_compaction)); - if (segment) { - increment.SetRealtimeOffsetRange(sealed_range.value()); + if (!segments.empty()) { + increment.SetRealtimeOffsetRange(OffsetRange(segments.front()->GetOffsetRange().begin, + segments.back()->GetOffsetRange().end)); } return increment; } @@ -242,15 +251,16 @@ Status RealtimePrimaryKeyWriter::FlushSegment( realtime_store_->CreateCommitReaders(segment)); PAIMON_ASSIGN_OR_RAISE( std::vector> realtime_primary_key_readers, - RealtimePrimaryKeyReaderFactory::CreateForCommit(std::move(readers), key_schema_, - write_schema_, memory_pool_)); + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(readers), key_schema_, schema_layout_->UserSchema(), memory_pool_)); std::vector> sorted_readers; sorted_readers.reserve(realtime_primary_key_readers.size()); for (std::unique_ptr& realtime_primary_key_reader : realtime_primary_key_readers) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, - PrimaryKeyTableUtils::CreateMergeFunction( - write_schema_, trimmed_primary_keys_, options_, memory_pool_)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + schema_layout_->UserSchema(), trimmed_primary_keys_, options_, memory_pool_)); sorted_readers.push_back(std::make_unique( std::move(realtime_primary_key_reader), key_comparator_, std::make_shared(std::move(merge_function)))); @@ -276,6 +286,10 @@ Status RealtimePrimaryKeyWriter::Sync() { Status RealtimePrimaryKeyWriter::Close() { return merge_tree_writer_->Close(); } +bool RealtimePrimaryKeyWriter::HasUnpreparedRealtimeData() const { + std::lock_guard lock(realtime_store_mutex_); + return has_building_data_ || !sealed_segments_.empty(); +} 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 86b624014..d3fb718bc 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -27,6 +27,7 @@ #include #include "paimon/core/core_options.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/utils/batch_writer.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" @@ -48,8 +49,7 @@ 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& transport_schema, + const std::shared_ptr& schema_layout, const std::vector& trimmed_primary_keys, const std::shared_ptr& key_comparator, const CoreOptions& options, const std::shared_ptr& realtime_context, @@ -58,6 +58,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { const std::shared_ptr& memory_pool); Status Write(std::unique_ptr&& batch) override; + Status Seal() override; Result PrepareCommit(bool wait_compaction) override; Status Compact(bool full_compaction) override; uint64_t GetMemoryUsage() const override; @@ -65,6 +66,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { Result CompactNotCompleted() override; Status Sync() override; Status Close() override; + bool HasUnpreparedRealtimeData() const override; std::shared_ptr GetMetrics() const override; private: @@ -72,8 +74,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { 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& transport_schema, + const std::shared_ptr& schema_layout, const std::shared_ptr& key_schema, const std::vector& trimmed_primary_keys, const std::shared_ptr& key_comparator, @@ -82,6 +83,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { const std::shared_ptr& memory_pool); Status FlushSegment(const std::shared_ptr& segment); + Status SealCurrentSegment(); std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; @@ -89,16 +91,16 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { std::shared_ptr merge_tree_writer_; std::shared_ptr realtime_context_; RealtimePartitionBucket partition_bucket_; - std::shared_ptr write_schema_; - std::shared_ptr realtime_input_schema_; - std::shared_ptr transport_schema_; + std::shared_ptr schema_layout_; + std::vector> sealed_segments_; std::shared_ptr key_schema_; std::vector trimmed_primary_keys_; std::shared_ptr key_comparator_; CoreOptions options_; int64_t next_offset_; int64_t last_sequence_number_; - std::mutex realtime_store_mutex_; + bool has_building_data_ = false; + mutable std::mutex realtime_store_mutex_; std::mutex prepare_mutex_; }; diff --git a/src/paimon/core/realtime/realtime_schema_layout.cpp b/src/paimon/core/realtime/realtime_schema_layout.cpp new file mode 100644 index 000000000..b67747f81 --- /dev/null +++ b/src/paimon/core/realtime/realtime_schema_layout.cpp @@ -0,0 +1,77 @@ +/* + * 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_schema_layout.h" + +#include + +#include "arrow/api.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +std::shared_ptr ValueKindField() { + return DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false); +} + +std::shared_ptr RealtimeOffsetField() { + return DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()); +} + +std::shared_ptr Prepend(const arrow::FieldVector& prefix, + const std::shared_ptr& schema) { + arrow::FieldVector fields = prefix; + fields.insert(fields.end(), schema->fields().begin(), schema->fields().end()); + return arrow::schema(std::move(fields), schema->metadata()); +} + +} // namespace + +Result> RealtimeSchemaLayout::Create( + RealtimeStoreMode mode, const std::shared_ptr& user_schema) { + if (!user_schema) { + return Status::Invalid("real-time schema layout requires a user schema"); + } + if (mode != RealtimeStoreMode::APPEND_ONLY && mode != RealtimeStoreMode::PRIMARY_KEY) { + return Status::Invalid("unknown real-time store mode"); + } + return std::unique_ptr(new RealtimeSchemaLayout(mode, user_schema)); +} + +RealtimeSchemaLayout::RealtimeSchemaLayout(RealtimeStoreMode mode, + const std::shared_ptr& user_schema) + : user_schema_(user_schema), input_schema_(Prepend({RealtimeOffsetField()}, user_schema_)) { + if (mode == RealtimeStoreMode::APPEND_ONLY) { + store_write_schema_ = input_schema_; + store_commit_schema_ = store_write_schema_; + commit_schema_ = user_schema_; + query_schema_ = Prepend({ValueKindField()}, user_schema_); + return; + } + + store_write_schema_ = SpecialFields::CompleteSequenceAndValueKindField(input_schema_); + store_commit_schema_ = store_write_schema_; + commit_schema_ = SpecialFields::CompleteSequenceAndValueKindField(user_schema_); + query_schema_ = commit_schema_; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_schema_layout.h b/src/paimon/core/realtime/realtime_schema_layout.h new file mode 100644 index 000000000..11668a30e --- /dev/null +++ b/src/paimon/core/realtime/realtime_schema_layout.h @@ -0,0 +1,88 @@ +/* + * 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 "arrow/type_fwd.h" +#include "paimon/realtime/realtime_store.h" +#include "paimon/result.h" + +namespace paimon { + +/// Defines the Arrow schemas at each real-time read and write boundary. +class RealtimeSchemaLayout { + public: + /// Creates all real-time boundary schemas for the requested store mode. + /// User fields retain their original order after any internal field prefix. + static Result> Create( + RealtimeStoreMode mode, const std::shared_ptr& user_schema); + + RealtimeSchemaLayout(const RealtimeSchemaLayout&) = delete; + RealtimeSchemaLayout& operator=(const RealtimeSchemaLayout&) = delete; + + /// Append-only fields: [user fields]. + /// Primary-key fields: [user fields]. + const std::shared_ptr& UserSchema() const { + return user_schema_; + } + + /// Append-only fields: [_REALTIME_OFFSET, user fields]. + /// Primary-key fields: [_REALTIME_OFFSET, user fields]. + const std::shared_ptr& InputSchema() const { + return input_schema_; + } + + /// Append-only fields: [_REALTIME_OFFSET, user fields]. + /// Primary-key fields: [_SEQUENCE_NUMBER, _VALUE_KIND, _REALTIME_OFFSET, user fields]. + const std::shared_ptr& StoreWriteSchema() const { + return store_write_schema_; + } + + /// Append-only fields: [_REALTIME_OFFSET, user fields]. + /// Primary-key fields: [_SEQUENCE_NUMBER, _VALUE_KIND, _REALTIME_OFFSET, user fields]. + const std::shared_ptr& StoreCommitSchema() const { + return store_commit_schema_; + } + + /// Append-only fields: [user fields]. + /// Primary-key fields: [_SEQUENCE_NUMBER, _VALUE_KIND, user fields]. + const std::shared_ptr& CommitSchema() const { + return commit_schema_; + } + + /// Append-only fields: [_VALUE_KIND, user fields]. + /// Primary-key fields: [_SEQUENCE_NUMBER, _VALUE_KIND, user fields]. + const std::shared_ptr& QuerySchema() const { + return query_schema_; + } + + private: + RealtimeSchemaLayout(RealtimeStoreMode mode, const std::shared_ptr& user_schema); + + std::shared_ptr user_schema_; + std::shared_ptr input_schema_; + std::shared_ptr store_write_schema_; + std::shared_ptr store_commit_schema_; + std::shared_ptr commit_schema_; + std::shared_ptr query_schema_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_schema_layout_test.cpp b/src/paimon/core/realtime/realtime_schema_layout_test.cpp new file mode 100644 index 000000000..397102b73 --- /dev/null +++ b/src/paimon/core/realtime/realtime_schema_layout_test.cpp @@ -0,0 +1,60 @@ +/* + * 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_schema_layout.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(RealtimeSchemaLayoutTest, TestSchemaLayouts) { + arrow::FieldVector value_fields = {arrow::field("key", arrow::int64(), false), + arrow::field("value", arrow::utf8())}; + std::shared_ptr user_schema = arrow::schema(value_fields); + ASSERT_OK_AND_ASSIGN(std::unique_ptr append_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, user_schema)); + ASSERT_EQ((std::vector{"_REALTIME_OFFSET", "key", "value"}), + append_layout->InputSchema()->field_names()); + ASSERT_TRUE(append_layout->InputSchema()->Equals(*append_layout->StoreWriteSchema())); + ASSERT_TRUE(append_layout->StoreWriteSchema()->Equals(*append_layout->StoreCommitSchema())); + ASSERT_TRUE(user_schema->Equals(*append_layout->CommitSchema())); + ASSERT_EQ((std::vector{"_VALUE_KIND", "key", "value"}), + append_layout->QuerySchema()->field_names()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr primary_key_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, user_schema)); + ASSERT_EQ(append_layout->InputSchema()->field_names(), + primary_key_layout->InputSchema()->field_names()); + ASSERT_EQ((std::vector{"_SEQUENCE_NUMBER", "_VALUE_KIND", "_REALTIME_OFFSET", + "key", "value"}), + primary_key_layout->StoreWriteSchema()->field_names()); + ASSERT_TRUE( + primary_key_layout->StoreWriteSchema()->Equals(*primary_key_layout->StoreCommitSchema())); + ASSERT_EQ((std::vector{"_SEQUENCE_NUMBER", "_VALUE_KIND", "key", "value"}), + primary_key_layout->CommitSchema()->field_names()); + ASSERT_TRUE(primary_key_layout->CommitSchema()->Equals(*primary_key_layout->QuerySchema())); +} + +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_store_read_pipeline.cpp b/src/paimon/core/realtime/realtime_store_read_pipeline.cpp index 4fe721a8e..0aaccb012 100644 --- a/src/paimon/core/realtime/realtime_store_read_pipeline.cpp +++ b/src/paimon/core/realtime/realtime_store_read_pipeline.cpp @@ -32,6 +32,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/realtime/realtime_offset_batch_reader.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/utils/nested_projection_utils.h" namespace paimon { @@ -87,18 +88,18 @@ class PhysicalToLogicalBatchReader : public BatchReader { "real-time physical-to-logical conversion requires a StructArray"); } auto struct_array = checked_pointer_cast(array); - const std::string value_kind_name = SpecialFields::ValueKind().Name(); - if (struct_array->num_fields() == 0 || - struct_array->struct_type()->field(0)->name() != value_kind_name) { - return Status::Invalid("real-time query batch must start with _VALUE_KIND"); + arrow::ArrayVector result_arrays; + arrow::FieldVector result_fields; + for (int32_t i = 0; i < struct_array->num_fields(); ++i) { + const std::shared_ptr& source_field = + struct_array->struct_type()->field(i); + if (SpecialFields::IsSystemField(source_field->name())) { + result_arrays.push_back(struct_array->field(i)); + result_fields.push_back(source_field); + } } - - arrow::ArrayVector result_arrays = {struct_array->field(0)}; - arrow::FieldVector result_fields = {struct_array->struct_type()->field(0)}; - result_arrays.reserve(logical_schema_->num_fields() + 1); - result_fields.reserve(logical_schema_->num_fields() + 1); for (const std::shared_ptr& read_field : logical_schema_->fields()) { - if (read_field->name() == value_kind_name) { + if (SpecialFields::IsSystemField(read_field->name())) { continue; } int32_t source_index = struct_array->struct_type()->GetFieldIndex(read_field->name()); @@ -139,13 +140,13 @@ class PhysicalToLogicalBatchReader : public BatchReader { } // namespace Result> RealtimeStoreReadPipeline::Create( - const std::shared_ptr& logical_schema, - const std::shared_ptr& realtime_write_schema, + const std::shared_ptr& logical_schema, const RealtimeSchemaLayout& schema_layout, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool) { - if (!logical_schema || !realtime_write_schema || !memory_pool || !arrow_pool) { + if (!logical_schema || !memory_pool || !arrow_pool) { return Status::Invalid("real-time store read pipeline requires schemas and memory pools"); } + const std::shared_ptr& store_write_schema = schema_layout.StoreWriteSchema(); std::map> plans; for (const std::shared_ptr& read_field : logical_schema->fields()) { @@ -153,7 +154,7 @@ Result> RealtimeStoreReadPipeline::Cr continue; } std::shared_ptr write_field = - realtime_write_schema->GetFieldByName(read_field->name()); + store_write_schema->GetFieldByName(read_field->name()); if (!write_field) { return Status::Invalid( fmt::format("selected-key MAP field {} does not exist in real-time write schema", @@ -167,7 +168,7 @@ Result> RealtimeStoreReadPipeline::Cr std::map> variant_plans; PAIMON_ASSIGN_OR_RAISE(variant_plans, VariantShreddingReadPlanFactory::CreateReadPlans( - logical_schema, realtime_write_schema, memory_pool)); + logical_schema, store_write_schema, memory_pool)); for (auto& [field_name, plan] : variant_plans) { if (!plans.emplace(field_name, std::move(plan)).second) { return Status::Invalid( @@ -176,11 +177,11 @@ Result> RealtimeStoreReadPipeline::Cr } bool needs_conversion = !plans.empty(); - // PK: _VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, then requested physical fields. + // PK: _SEQUENCE_NUMBER, _VALUE_KIND, _REALTIME_OFFSET, then requested physical fields. // Append: _REALTIME_OFFSET, then requested physical fields. arrow::FieldVector store_read_fields; - store_read_fields.reserve(realtime_write_schema->num_fields()); - for (const std::shared_ptr& write_field : realtime_write_schema->fields()) { + store_read_fields.reserve(store_write_schema->num_fields()); + for (const std::shared_ptr& write_field : store_write_schema->fields()) { if (SpecialFields::IsSystemField(write_field->name())) { store_read_fields.push_back(write_field); } diff --git a/src/paimon/core/realtime/realtime_store_read_pipeline.h b/src/paimon/core/realtime/realtime_store_read_pipeline.h index c660b1b00..24fba4e58 100644 --- a/src/paimon/core/realtime/realtime_store_read_pipeline.h +++ b/src/paimon/core/realtime/realtime_store_read_pipeline.h @@ -36,16 +36,15 @@ class Schema; namespace paimon { +class RealtimeSchemaLayout; + /// Builds the schema requested from a `RealtimeStore` and converts a store reader into the /// logical representation expected by table read. class RealtimeStoreReadPipeline { public: - /// `realtime_write_schema` is the complete schema written to `RealtimeStore`, including its - /// system fields. static Result> Create( const std::shared_ptr& logical_schema, - const std::shared_ptr& realtime_write_schema, - const std::shared_ptr& memory_pool, + const RealtimeSchemaLayout& schema_layout, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool); const std::shared_ptr& StoreReadSchema() const { diff --git a/src/paimon/core/realtime/realtime_store_read_pipeline_test.cpp b/src/paimon/core/realtime/realtime_store_read_pipeline_test.cpp index 791b75740..60e9c84c3 100644 --- a/src/paimon/core/realtime/realtime_store_read_pipeline_test.cpp +++ b/src/paimon/core/realtime/realtime_store_read_pipeline_test.cpp @@ -30,6 +30,7 @@ #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" #include "paimon/data/variant.h" #include "paimon/memory/memory_pool.h" @@ -74,7 +75,9 @@ TEST(RealtimeStoreReadPipelineTest, SelectedMapKeysAsMapAndStruct) { DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()); auto id_field = arrow::field("id", arrow::int64()); auto tags_field = arrow::field("tags", map_type); - auto write_schema = arrow::schema({offset_field, id_field, tags_field}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, + arrow::schema({id_field, tags_field}))); std::shared_ptr source_type = arrow::struct_({value_kind_field, offset_field, id_field, tags_field}); std::shared_ptr source = @@ -88,8 +91,8 @@ TEST(RealtimeStoreReadPipelineTest, SelectedMapKeysAsMapAndStruct) { auto selected_map = MapReadField(map_type, "c,a,missing"); ASSERT_OK_AND_ASSIGN( std::unique_ptr map_pipeline, - RealtimeStoreReadPipeline::Create(arrow::schema({value_kind_field, id_field, selected_map}), - write_schema, pool, GetArrowPool(pool))); + RealtimeStoreReadPipeline::Create(arrow::schema({id_field, selected_map}), *schema_layout, + pool, GetArrowPool(pool))); ASSERT_TRUE(map_pipeline->StoreReadSchema()->field(2)->type()->Equals(map_type)); auto map_source_reader = std::make_unique( source, source_type, static_cast(source->length())); @@ -113,10 +116,10 @@ TEST(RealtimeStoreReadPipelineTest, SelectedMapKeysAsMapAndStruct) { ASSERT_OK_AND_ASSIGN(std::shared_ptr selected_struct, MapAccessField(tags_field, {"a", "missing"})); - ASSERT_OK_AND_ASSIGN(std::unique_ptr struct_pipeline, - RealtimeStoreReadPipeline::Create( - arrow::schema({value_kind_field, id_field, selected_struct}), - write_schema, pool, GetArrowPool(pool))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr struct_pipeline, + RealtimeStoreReadPipeline::Create(arrow::schema({id_field, selected_struct}), + *schema_layout, pool, GetArrowPool(pool))); ASSERT_EQ(struct_pipeline->StoreReadSchema()->field(2)->type()->id(), arrow::Type::MAP); auto struct_source_reader = std::make_unique( source, source_type, static_cast(source->length())); @@ -148,7 +151,9 @@ TEST(RealtimeStoreReadPipelineTest, VariantAccessOnLogicalVariant) { DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()); auto id_field = arrow::field("id", arrow::int32()); auto variant_field = VariantTypeUtils::ToArrowField("v"); - auto write_schema = arrow::schema({offset_field, id_field, variant_field}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, + arrow::schema({id_field, variant_field}))); ASSERT_OK_AND_ASSIGN(std::shared_ptr data, VariantTestData::BuildVariantBatch(id_field, variant_field, {R"({"a":5,"city":"hangzhou"})"}, pool, @@ -173,11 +178,11 @@ TEST(RealtimeStoreReadPipelineTest, VariantAccessOnLogicalVariant) { auto access_field_result = arrow::ImportField(c_access_field.get()); ASSERT_TRUE(access_field_result.ok()) << access_field_result.status().ToString(); std::shared_ptr access_field = access_field_result.ValueOrDie(); - auto read_schema = arrow::schema({value_kind_field, id_field, access_field}); + auto read_schema = arrow::schema({id_field, access_field}); ASSERT_OK_AND_ASSIGN( std::unique_ptr pipeline, - RealtimeStoreReadPipeline::Create(read_schema, write_schema, pool, GetArrowPool(pool))); + RealtimeStoreReadPipeline::Create(read_schema, *schema_layout, pool, GetArrowPool(pool))); ASSERT_TRUE(pipeline->StoreReadSchema()->field(2)->type()->Equals(variant_field->type())); auto source_reader = std::make_unique( source, source->type(), static_cast(source->length())); 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 ea93d4a53..d5151ed43 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -28,6 +28,7 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/reader/predicate_batch_reader.h" #include "paimon/common/types/data_field.h" @@ -38,8 +39,8 @@ #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/operation/raw_file_split_read.h" #include "paimon/core/realtime/realtime_context_impl.h" -#include "paimon/core/realtime/realtime_offset_utils.h" #include "paimon/core/realtime/realtime_reader.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/realtime/realtime_store_read_pipeline.h" #include "paimon/core/table/source/append_count_reader.h" #include "paimon/core/table/source/realtime_split.h" @@ -161,12 +162,13 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( std::shared_ptr table_schema = DataField::ConvertDataFieldsToArrowSchema(context_->GetTableSchema()->Fields()); - std::shared_ptr realtime_write_schema = - RealtimeOffsetUtils::CreateInputSchema(table_schema); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, table_schema)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr pipeline, RealtimeStoreReadPipeline::Create( - context_->GetReadSchema(), realtime_write_schema, - context_->GetMemoryPool(), context_->GetArrowMemoryPool())); + context_->GetReadSchema(), *schema_layout, context_->GetMemoryPool(), + context_->GetArrowMemoryPool())); const std::shared_ptr& store_read_schema = pipeline->StoreReadSchema(); auto c_read_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*store_read_schema, c_read_schema.get())); @@ -196,6 +198,8 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( pipeline->Wrap(std::move(memory_reader), OffsetRange(realtime_split->CommittedEndOffset(), realtime_split->MemoryEndOffset()))); + memory_reader = std::make_unique( + std::move(memory_reader), context_->GetArrowMemoryPool()); if (context_->EnablePredicateFilter() && context_->GetPredicate()) { PAIMON_ASSIGN_OR_RAISE( memory_reader, diff --git a/src/paimon/core/table/source/data_table_batch_scan.cpp b/src/paimon/core/table/source/data_table_batch_scan.cpp index 0414045ba..924f3fd5c 100644 --- a/src/paimon/core/table/source/data_table_batch_scan.cpp +++ b/src/paimon/core/table/source/data_table_batch_scan.cpp @@ -36,7 +36,8 @@ class DataSplit; DataTableBatchScan::DataTableBatchScan(bool pk_table, const CoreOptions& core_options, const std::shared_ptr& snapshot_reader, - bool read_optimized, std::optional push_down_limit) + bool read_optimized, std::optional push_down_limit, + bool realtime_pk_scan) : AbstractTableScan(core_options, snapshot_reader), push_down_limit_(push_down_limit), logger_(Logger::GetLogger("DataTableBatchScan")) { @@ -47,9 +48,14 @@ DataTableBatchScan::DataTableBatchScan(bool pk_table, const CoreOptions& core_op snapshot_reader_->EnableValueFilter(); } else if (pk_table && (core_options.DeletionVectorsEnabled() || core_options.GetMergeEngine() == MergeEngine::FIRST_ROW)) { - auto level_filter = [](int32_t level) -> bool { return level > 0; }; - snapshot_reader_->WithLevelFilter(level_filter); - snapshot_reader_->EnableValueFilter(); + if (realtime_pk_scan) { + snapshot_reader_->EnableValueFilterForLevels( + [](int32_t level) -> bool { return level > 0; }); + } else { + auto level_filter = [](int32_t level) -> bool { return level > 0; }; + snapshot_reader_->WithLevelFilter(level_filter); + snapshot_reader_->EnableValueFilter(); + } } if (core_options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { snapshot_reader_->OnlyReadRealBuckets(); diff --git a/src/paimon/core/table/source/data_table_batch_scan.h b/src/paimon/core/table/source/data_table_batch_scan.h index 6f9c2246e..4f0ef65e3 100644 --- a/src/paimon/core/table/source/data_table_batch_scan.h +++ b/src/paimon/core/table/source/data_table_batch_scan.h @@ -37,7 +37,7 @@ class DataTableBatchScan : public AbstractTableScan { public: DataTableBatchScan(bool pk_table, const CoreOptions& core_options, const std::shared_ptr& snapshot_reader, bool read_optimized, - std::optional push_down_limit); + std::optional push_down_limit, bool realtime_pk_scan); Result> CreatePlan() override; 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 9e81066a8..c25011bcf 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -37,9 +37,9 @@ #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" #include "paimon/core/realtime/realtime_context_impl.h" -#include "paimon/core/realtime/realtime_offset_utils.h" #include "paimon/core/realtime/realtime_primary_key_reader.h" #include "paimon/core/realtime/realtime_reader.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/realtime/realtime_store_read_pipeline.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" @@ -61,22 +61,26 @@ namespace { Result> CreateRealtimePrimaryKeyLogicalSchema( const std::shared_ptr& key_schema, const std::shared_ptr& value_schema) { - arrow::FieldVector transport_value_fields; - transport_value_fields.reserve(key_schema->num_fields() + value_schema->num_fields()); + arrow::FieldVector query_value_fields; + query_value_fields.reserve(key_schema->num_fields() + value_schema->num_fields()); std::unordered_set field_ids; for (const std::shared_ptr& field : key_schema->fields()) { PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); if (field_ids.insert(field_id).second) { - transport_value_fields.push_back(field); + query_value_fields.push_back(field); } } for (const std::shared_ptr& field : value_schema->fields()) { PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); if (field_ids.insert(field_id).second) { - transport_value_fields.push_back(field); + query_value_fields.push_back(field); } } - return RealtimePrimaryKeyLayout::CreateSchema(transport_value_fields); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, + arrow::schema(std::move(query_value_fields)))); + return schema_layout->QuerySchema(); } Result>> CreateMemoryReaders( @@ -89,13 +93,12 @@ Result>> CreateMemoryReaders( const std::shared_ptr& memory_pool) { std::shared_ptr table_write_schema = DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields()); - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(table_write_schema); - std::shared_ptr realtime_write_schema = - RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, table_write_schema)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr pipeline, - RealtimeStoreReadPipeline::Create(logical_schema, realtime_write_schema, memory_pool, + RealtimeStoreReadPipeline::Create(logical_schema, *schema_layout, memory_pool, context->GetArrowMemoryPool())); const std::shared_ptr& store_read_schema = pipeline->StoreReadSchema(); auto c_schema = std::make_unique(); diff --git a/src/paimon/core/table/source/merge_tree_split_generator.cpp b/src/paimon/core/table/source/merge_tree_split_generator.cpp index 925ce2269..5e2caea5d 100644 --- a/src/paimon/core/table/source/merge_tree_split_generator.cpp +++ b/src/paimon/core/table/source/merge_tree_split_generator.cpp @@ -105,7 +105,21 @@ Result> MergeTreeSplitGenerator::SplitFo std::vector split_groups; split_groups.reserve(metas_vec.size()); for (auto& metas : metas_vec) { - if (metas.size() == 1 && WithoutDeleteRow(metas[0])) { + bool packed_raw_convertible = true; + std::set packed_levels; + for (const std::shared_ptr& meta : metas) { + if (meta->level == 0 || !WithoutDeleteRow(meta)) { + packed_raw_convertible = false; + } + packed_levels.insert(meta->level); + } + const bool single_file_raw_convertible = + metas.size() == 1 && WithoutDeleteRow(metas[0]) && + (!deletion_vectors_enabled_ || metas[0]->level > 0); + if (single_file_raw_convertible || + (packed_raw_convertible && + (deletion_vectors_enabled_ || merge_engine_ == MergeEngine::FIRST_ROW || + packed_levels.size() == 1))) { split_groups.push_back( SplitGenerator::SplitGroup::RawConvertibleGroup(std::move(metas))); } else { diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index 4c3968dc3..5cfb26c6a 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -25,6 +25,7 @@ #include #include +#include "fmt/format.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/snapshot.h" @@ -78,6 +79,31 @@ Result RealtimeTableScan::LoadCommittedOffsets( file_system_); } +Status RealtimeTableScan::ValidateSnapshotProgress( + const std::vector& memory_views, + const RealtimeOffsetMap& snapshot_offsets, const RealtimeOffsetMap& context_offsets) const { + for (const RealtimePartitionBucketView& memory : memory_views) { + if (scan_filter_->GetBucketFilter() && + memory.partition_bucket.bucket != scan_filter_->GetBucketFilter().value()) { + continue; + } + if (!MatchPartition(memory.partition_bucket.partition)) { + continue; + } + const int64_t snapshot_offset = + GetCommittedEndOffset(snapshot_offsets, memory.partition_bucket); + const int64_t context_offset = + GetCommittedEndOffset(context_offsets, memory.partition_bucket); + if (snapshot_offset < context_offset) { + return Status::Invalid(fmt::format( + "real-time scan snapshot offset {} is behind context committed offset {} for " + "bucket {}; the missing range may already have been reclaimed", + snapshot_offset, context_offset, memory.partition_bucket.bucket)); + } + } + return Status::OK(); +} + Result RealtimeTableScan::CollectActiveMemoryViews( std::vector&& memory_views, const RealtimeOffsetMap& committed_offsets) const { @@ -196,12 +222,14 @@ RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, boo Result> RealtimeTableScan::CreatePlan() { // Memory is pinned first. If a commit and reclaim happens before disk planning, the old memory // remains alive in these views and the selected snapshot offset removes its covered prefix. - PAIMON_ASSIGN_OR_RAISE(std::vector memory_views, - realtime_context_->AcquireReadViews()); + PAIMON_ASSIGN_OR_RAISE(RealtimeReadState read_state, realtime_context_->AcquireReadState()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr disk_plan, disk_scan_->CreatePlan()); PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap committed_offsets, LoadCommittedOffsets(disk_plan)); - PAIMON_ASSIGN_OR_RAISE(MemoryViewMap active_memory, - CollectActiveMemoryViews(std::move(memory_views), committed_offsets)); + PAIMON_RETURN_NOT_OK(ValidateSnapshotProgress(read_state.views, committed_offsets, + read_state.committed_offsets)); + PAIMON_ASSIGN_OR_RAISE( + MemoryViewMap active_memory, + CollectActiveMemoryViews(std::move(read_state.views), committed_offsets)); PAIMON_ASSIGN_OR_RAISE(std::vector> splits, CreateRealtimeSplits(disk_plan->Splits(), std::move(active_memory), committed_offsets, disk_plan->SnapshotId())); diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h index 6a4ad3e91..821e67fd7 100644 --- a/src/paimon/core/table/source/realtime_table_scan.h +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -61,6 +61,10 @@ class RealtimeTableScan : public TableScan { Result LoadCommittedOffsets(const std::shared_ptr& disk_plan) const; + Status ValidateSnapshotProgress(const std::vector& memory_views, + const RealtimeOffsetMap& snapshot_offsets, + const RealtimeOffsetMap& context_offsets) const; + Result CollectActiveMemoryViews( std::vector&& memory_views, const RealtimeOffsetMap& committed_offsets) const; diff --git a/src/paimon/core/table/source/snapshot/snapshot_reader.h b/src/paimon/core/table/source/snapshot/snapshot_reader.h index 35134d22a..47860d4bf 100644 --- a/src/paimon/core/table/source/snapshot/snapshot_reader.h +++ b/src/paimon/core/table/source/snapshot/snapshot_reader.h @@ -78,6 +78,11 @@ class SnapshotReader { return this; } + SnapshotReader* EnableValueFilterForLevels(const std::function& level_filter) { + scan_->EnableValueFilterForLevels(level_filter); + return this; + } + SnapshotReader* OnlyReadRealBuckets() { scan_->OnlyReadRealBuckets(); return this; diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index e4db33fb4..4e3d911f4 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -255,8 +255,12 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& return Status::Invalid("real-time union read does not support global index splits"); } StartupMode startup_mode = core_options.GetStartupMode(); - if (!(startup_mode == StartupMode::LatestFull() || startup_mode == StartupMode::Latest())) { - return Status::Invalid("real-time union read requires the latest snapshot"); + const bool specified_snapshot = + startup_mode == StartupMode::FromSnapshot() && core_options.GetScanSnapshotId().has_value(); + if (!(startup_mode == StartupMode::LatestFull() || startup_mode == StartupMode::Latest() || + specified_snapshot)) { + return Status::Invalid( + "real-time union read requires the latest snapshot or an explicit snapshot id"); } return Status::OK(); } @@ -350,7 +354,8 @@ Result> NewDataTableScan(const std::shared_ptr(core_options, snapshot_reader); } auto batch_scan = std::make_unique( - /*pk_table=*/pk_table, core_options, snapshot_reader, read_optimized, context->GetLimit()); + /*pk_table=*/pk_table, core_options, snapshot_reader, read_optimized, context->GetLimit(), + /*realtime_pk_scan=*/pk_table && context->GetRealtimeContext() != nullptr); if (context->GetRealtimeContext()) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context, RealtimeContextImpl::Cast(context->GetRealtimeContext())); diff --git a/src/paimon/core/utils/batch_writer.h b/src/paimon/core/utils/batch_writer.h index 6c31b4868..abe4b4c4a 100644 --- a/src/paimon/core/utils/batch_writer.h +++ b/src/paimon/core/utils/batch_writer.h @@ -47,6 +47,11 @@ class BatchWriter { /// Add a record batch to the writer. virtual Status Write(std::unique_ptr&& batch) = 0; + /// Seal the current in-memory segment without flushing it. + virtual Status Seal() { + return Status::NotImplemented("seal is not supported by this batch writer"); + } + /// Compact files related to the writer. Note that compaction process is only submitted and may /// not be completed when the method returns. /// @@ -70,6 +75,11 @@ class BatchWriter { /// Close this writer, the call will delete newly generated but not committed files. virtual Status Close() = 0; + /// Whether this writer owns real-time data not covered by a successful prepare. + virtual bool HasUnpreparedRealtimeData() const { + return false; + } + virtual std::shared_ptr GetMetrics() const = 0; }; diff --git a/src/paimon/core/utils/primary_key_table_utils.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index 2ca444a7a..639bb8c82 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -122,10 +122,7 @@ Status PrimaryKeyTableUtils::ValidateRealtimeOptions(const CoreOptions& options, if (options.GetChangelogProducer() != ChangelogProducer::NONE) { return Status::NotImplemented("PK realtime supports only the NONE changelog producer"); } - if (options.DeletionVectorsEnabled()) { - return Status::NotImplemented("PK realtime does not support deletion vectors"); - } - if (options.NeedLookup()) { + if (options.NeedLookup() && !options.DeletionVectorsEnabled()) { return Status::NotImplemented("PK realtime does not support lookup"); } PAIMON_ASSIGN_OR_RAISE(std::vector primary_key_fields, 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 796922336..5f8a7cb3a 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -58,7 +58,9 @@ std::shared_ptr PkSchema( } // namespace TEST(PrimaryKeyTableUtilsTest, TestSupportedRealtimeOptions) { - ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::BUCKET, "1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}})); ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema())); } @@ -114,8 +116,6 @@ TEST(PrimaryKeyTableUtilsTest, TestRealtimeReportsSpecificLookupErrors) { const std::vector, std::string>> cases = { {{{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, "PK realtime does not support lookup"}, - {{{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, - "PK realtime does not support deletion vectors"}, {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, "PK realtime supports only the NONE changelog producer"}, {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "lookup"}}, diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 95adf74fc..73fa1b276 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -52,11 +52,11 @@ #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/realtime/realtime_context_impl.h" -#include "paimon/core/realtime/realtime_offset_utils.h" -#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/realtime/realtime_schema_layout.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/table/sink/commit_message_impl.h" +#include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/realtime_split.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" @@ -431,8 +431,10 @@ class RealtimeWriteInteTest : public ::testing::Test { } json += "]"; - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, schema_)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ipc::internal::json::ArrayFromJSON( arrow::struct_(realtime_schema->fields()), json)); @@ -490,8 +492,10 @@ class RealtimeWriteInteTest : public ::testing::Test { } json += "]"; - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, schema_)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ipc::internal::json::ArrayFromJSON( arrow::struct_(realtime_schema->fields()), json)); @@ -502,8 +506,10 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> MakeUnpartitionedBatchFromJson( const std::string& json) const { - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, schema_)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ipc::internal::json::ArrayFromJSON( arrow::struct_(realtime_schema->fields()), json)); @@ -590,14 +596,15 @@ class RealtimeWriteInteTest : public ::testing::Test { } Result CompactAndCommit(const std::map& partition, - int32_t bucket, int64_t commit_identifier) const { + int32_t bucket, int64_t commit_identifier, + bool full_compaction = true) const { WriteContextBuilder write_builder(table_path_, commit_user_); - write_builder.SetOptions(options_).WithStreamingMode(true); + write_builder.SetOptions(options_).WithStreamingMode(true).WithTempDirectory( + PathUtil::JoinPath(dir_->Str(), "compact-tmp")); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr write_context, write_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr compaction_writer, FileStoreWrite::Create(std::move(write_context))); - PAIMON_RETURN_NOT_OK(compaction_writer->Compact(partition, bucket, - /*full_compaction=*/true)); + PAIMON_RETURN_NOT_OK(compaction_writer->Compact(partition, bucket, full_compaction)); PAIMON_ASSIGN_OR_RAISE( std::vector> compaction_messages, compaction_writer->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); @@ -729,10 +736,17 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> ReadRows( const std::shared_ptr& plan, const std::shared_ptr& realtime_context) const { + return ReadRows(plan, realtime_context, /*predicate=*/nullptr, + /*enable_predicate_filter=*/false); + } + + Result> ReadRows(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context, + const std::shared_ptr& predicate, + bool enable_predicate_filter) const { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_result, - ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, - /*predicate=*/nullptr, - /*enable_predicate_filter=*/false)); + ReadPlan(plan, realtime_context, {"id", "payload", "pt"}, predicate, + enable_predicate_filter)); const std::shared_ptr& result = read_result; std::vector rows; @@ -790,10 +804,10 @@ class RealtimeWriteInteTest : public ::testing::Test { 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()); + PAIMON_ASSIGN_OR_RAISE(RealtimeReadState read_state, + realtime_context_impl->AcquireReadState()); uint64_t memory_usage = 0; - for (const RealtimePartitionBucketView& view : views) { + for (const RealtimePartitionBucketView& view : read_state.views) { memory_usage += view.store->GetMemoryUsage(); } return memory_usage; @@ -803,8 +817,9 @@ class RealtimeWriteInteTest : public ::testing::Test { 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()); + PAIMON_ASSIGN_OR_RAISE(RealtimeReadState read_state, + realtime_context_impl->AcquireReadState()); + const std::vector& views = read_state.views; if (views.size() != 1) { return Status::Invalid("expected one PK real-time read view"); } @@ -818,11 +833,11 @@ class RealtimeWriteInteTest : public ::testing::Test { auto read_schema = std::make_unique(); std::shared_ptr value_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema.value()->Fields()); - std::shared_ptr realtime_input_schema = - RealtimeOffsetUtils::CreateInputSchema(value_schema); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema( - *RealtimePrimaryKeyLayout::CreateSchema(realtime_input_schema->fields()), - read_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, value_schema)); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*schema_layout->StoreWriteSchema(), read_schema.get())); ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr}; PAIMON_ASSIGN_OR_RAISE( @@ -906,6 +921,13 @@ class RealtimeWriteInteTest : public ::testing::Test { return RealtimeCommitProperties::ReadOffsets(snapshot, options.GetFileSystem()); } + Status PrepareAndClose(FileStoreWrite* writer) const { + PAIMON_ASSIGN_OR_RAISE(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + static_cast(progress); + return writer->Close(); + } + void FinalizeCommitAndCheck(FileStoreWrite* writer, std::vector realtime_commits, int64_t prepare_identifier, std::vector expected_rows) const { @@ -1049,6 +1071,52 @@ TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } +TEST_F(RealtimeWriteInteTest, TestSealOnlySealsRealtimeSegment) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr 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(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(uint64_t memory_before_seal, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_GT(memory_before_seal, 0); + ASSERT_OK(writer->Seal()); + ASSERT_OK_AND_ASSIGN(uint64_t memory_after_seal, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(memory_before_seal, memory_after_seal); + ASSERT_OK_AND_ASSIGN(std::vector sealed_rows, ReadRows(realtime_context)); + ASSERT_EQ(first_rows, sealed_rows); + ASSERT_OK_AND_ASSIGN(std::vector disk_rows_before_prepare, ReadRows()); + ASSERT_TRUE(disk_rows_before_prepare.empty()); + + std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + std::vector expected_rows = first_rows; + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector rows_before_prepare, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows_before_prepare); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 5), progress[0].offset_range); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_prepare, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows_after_prepare); + 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 rows_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows_after_refresh); + ASSERT_OK(writer->Close()); + realtime_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); +} + TEST_F(RealtimeWriteInteTest, TestSparseExternalOffsetsCommitAndRecover) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -1081,6 +1149,51 @@ TEST_F(RealtimeWriteInteTest, TestSparseExternalOffsetsCommitAndRecover) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestExplicitSnapshotIdWithSparseTailAndReclaimBoundary) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ResetExternalOffset(/*partition=*/{}, /*bucket=*/0, /*next_offset=*/10); + const std::vector disk_rows = {{1, "disk", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + 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 disk_snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(disk_snapshot_id)); + + // Offset gaps are valid. The snapshot high-watermark is 11 while the next memory row starts at + // 30, so a real-time scan must not require the two numeric offsets to be adjacent. + ResetExternalOffset(/*partition=*/{}, /*bucket=*/0, /*next_offset=*/30); + const std::vector memory_rows = {{2, "memory", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + options_[Options::SCAN_SNAPSHOT_ID] = std::to_string(disk_snapshot_id); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(disk_snapshot_id, plan->SnapshotId()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "disk", "p0"}, {2, "memory", "p0"}}), actual_rows); + + options_.erase(Options::SCAN_SNAPSHOT_ID); + ASSERT_OK_AND_ASSIGN(std::vector memory_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t memory_snapshot_id, + Commit(memory_progress, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(memory_snapshot_id)); + + options_[Options::SCAN_SNAPSHOT_ID] = std::to_string(disk_snapshot_id); + ASSERT_NOK_WITH_MSG(CreatePlan(realtime_context, /*predicate=*/nullptr), + "behind context committed offset"); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPkRead) { CreatePkTable(); auto saw_query_predicate = std::make_shared>(false); @@ -1130,7 +1243,7 @@ TEST_F(RealtimeWriteInteTest, TestPkRead) { predicate, /*enable_predicate_filter=*/true)); ASSERT_EQ(nullptr, filtered_result); ASSERT_FALSE(saw_query_predicate->load(std::memory_order_acquire)); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); writer.reset(); ASSERT_OK_AND_ASSIGN(std::shared_ptr lifetime_plan, @@ -1304,7 +1417,7 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { .ValueOrDie(); ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) << result->ToString(); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPkMergeAllDiskSplitsWithMemory) { @@ -1354,7 +1467,7 @@ TEST_F(RealtimeWriteInteTest, TestPkMergeAllDiskSplitsWithMemory) { {20, "disk-20", "p0"}, {21, "disk-21", "p0"}}), actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { @@ -1370,8 +1483,9 @@ TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter(realtime_context)); - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::PRIMARY_KEY, schema_)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); auto make_batch = [&](const std::string& json) -> Result> { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ipc::internal::json::ArrayFromJSON( @@ -1438,7 +1552,7 @@ TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { .ValueOrDie(); ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) << actual->ToString(); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPkKeylessProjection) { @@ -1465,7 +1579,7 @@ TEST_F(RealtimeWriteInteTest, TestPkKeylessProjection) { arrow::schema({arrow::field("payload", arrow::utf8())}), R"([ [0, "memory"] ])"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestCompositePkKeylessProjection) { @@ -1492,7 +1606,7 @@ TEST_F(RealtimeWriteInteTest, TestCompositePkKeylessProjection) { arrow::schema({arrow::field("pt", arrow::utf8())}), R"([ [0, "memory"] ])"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { @@ -1552,7 +1666,7 @@ TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { {2, "b", "sealed-2b"}, {3, "c", "disk-3c"}}), actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { @@ -1997,7 +2111,7 @@ TEST_F(RealtimeWriteInteTest, TestAppendScanKeepsDiskSplitsIndependent) { ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { @@ -2160,6 +2274,203 @@ TEST_F(RealtimeWriteInteTest, TestRealtimeWriteAcrossAppendCompaction) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkDvPredicateAcrossHighLevelLevel0AndMemory) { + options_[Options::FILE_FORMAT] = "parquet"; + options_[Options::DELETION_VECTORS_ENABLED] = "true"; + 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::string matching_payload = "match-" + std::string(2048, 'X'); + const std::vector base_rows = {{1, matching_payload, "p0"}, + {2, matching_payload, "p0"}, + {3, matching_payload, "p0"}, + {4, "base-four", "p0"}, + {5, "base-five", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(base_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t base_snapshot_id, Commit(base_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(base_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(Snapshot full_compact_snapshot, + CompactAndCommit(/*partition=*/{}, /*bucket=*/0, + /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(full_compact_snapshot.Id())); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_full_compaction, ReadRows(realtime_context)); + ASSERT_EQ(base_rows, rows_after_full_compaction); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr dv_update_batch, + MakeBatch({Row{5, "dv-current-five", "p0"}}, /*partitioned=*/false, + /*bucket=*/0, {RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(dv_update_batch))); + ASSERT_OK_AND_ASSIGN(std::vector dv_update_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_OK_AND_ASSIGN(int64_t dv_update_snapshot_id, + Commit(dv_update_progress, /*commit_identifier=*/2)); + ASSERT_OK(writer->RefreshCommittedSnapshot(dv_update_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(Snapshot dv_compact_snapshot, + CompactAndCommit(/*partition=*/{}, /*bucket=*/0, + /*commit_identifier=*/3, + /*full_compaction=*/false)); + ASSERT_OK(writer->RefreshCommittedSnapshot(dv_compact_snapshot.Id())); + const std::vector expected_high_level_rows = {{1, matching_payload, "p0"}, + {2, matching_payload, "p0"}, + {3, matching_payload, "p0"}, + {4, "base-four", "p0"}, + {5, "dv-current-five", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector high_level_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_high_level_rows, high_level_rows); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr level0_batch, + MakeBatch({Row{1, "level0-current-one", "p0"}, Row{2, "deleted-two", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(level0_batch))); + ASSERT_OK_AND_ASSIGN(std::vector level0_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/4)); + ASSERT_OK_AND_ASSIGN(int64_t level0_snapshot_id, + Commit(level0_progress, /*commit_identifier=*/4)); + ASSERT_OK(writer->RefreshCommittedSnapshot(level0_snapshot_id)); + const std::vector expected_level0_rows = {{1, "level0-current-one", "p0"}, + {3, matching_payload, "p0"}, + {4, "base-four", "p0"}, + {5, "dv-current-five", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector level0_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_level0_rows, level0_rows); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr memory_batch, + MakeBatch({Row{3, "memory-current-three", "p0"}, Row{7, matching_payload, "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(memory_batch))); + const std::vector expected_memory_rows = {{1, "level0-current-one", "p0"}, + {3, "memory-current-three", "p0"}, + {4, "base-four", "p0"}, + {5, "dv-current-five", "p0"}, + {7, matching_payload, "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector memory_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_memory_rows, memory_rows); + + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, matching_payload.data(), matching_payload.size())); + options_[Options::SCAN_SNAPSHOT_ID] = std::to_string(level0_snapshot_id); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + ASSERT_EQ(level0_snapshot_id, plan->SnapshotId()); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, realtime_split); + ASSERT_GT(realtime_split->MemoryEndOffset(), realtime_split->CommittedEndOffset()); + + bool has_level0_file = false; + bool has_high_level_file = false; + bool has_high_level_deletion_vector = false; + for (const std::shared_ptr& disk_split : realtime_split->DiskSplits()) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + ASSERT_NE(nullptr, data_split); + const std::vector>& files = data_split->DataFiles(); + const std::vector>& deletion_files = + data_split->DeletionFiles(); + ASSERT_TRUE(deletion_files.empty() || deletion_files.size() == files.size()); + for (size_t i = 0; i < files.size(); ++i) { + if (files[i]->level == 0) { + has_level0_file = true; + } else { + has_high_level_file = true; + if (!deletion_files.empty() && deletion_files[i]) { + has_high_level_deletion_vector = true; + } + } + } + } + ASSERT_TRUE(has_level0_file); + ASSERT_TRUE(has_high_level_file); + ASSERT_TRUE(has_high_level_deletion_vector); + + // Predicate pushdown may return false-positive candidates, but a newer unfiltered L0 or + // memory record must still suppress every matching old version from the high levels. + ASSERT_OK_AND_ASSIGN(std::vector candidates, ReadRows(plan, realtime_context, predicate, + /*enable_predicate_filter=*/false)); + ASSERT_EQ((std::vector{{1, "level0-current-one", "p0"}, + {3, "memory-current-three", "p0"}, + {4, "base-four", "p0"}, + {7, matching_payload, "p0"}}), + candidates); + ASSERT_OK(PrepareAndClose(writer.get())); +} + +TEST_F(RealtimeWriteInteTest, TestPkDvMergesDuplicateKeysAcrossLevel0RunsAndMemory) { + options_[Options::FILE_FORMAT] = "parquet"; + options_[Options::DELETION_VECTORS_ENABLED] = "true"; + 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 first_level0_batch, + MakeBatch({Row{1, "level0-old-one", "p0"}, Row{2, "level0-old-two", "p0"}, + Row{3, "level0-stable-three", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_level0_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_level0_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t first_level0_snapshot_id, + Commit(first_level0_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(first_level0_snapshot_id)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr second_level0_batch, + MakeBatch({Row{1, "level0-new-one", "p0"}, Row{2, "level0-deleted-two", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(second_level0_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_level0_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t second_level0_snapshot_id, + Commit(second_level0_progress, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(second_level0_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch({Row{1, "memory-new-one", "p0"}}, /*partitioned=*/false, + /*bucket=*/0, {RecordBatch::RowKind::UPDATE_AFTER})); + 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); + size_t level0_file_count = 0; + for (const std::shared_ptr& disk_split : realtime_split->DiskSplits()) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + ASSERT_NE(nullptr, data_split); + for (const std::shared_ptr& file : data_split->DataFiles()) { + if (file->level == 0) { + ++level0_file_count; + } + } + } + ASSERT_GE(level0_file_count, 2); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "memory-new-one", "p0"}, {3, "level0-stable-three", "p0"}}), + actual_rows); + ASSERT_OK(PrepareAndClose(writer.get())); +} + TEST_F(RealtimeWriteInteTest, TestRealtimeOffsetFileLifecycle) { options_[Options::SNAPSHOT_NUM_RETAINED_MIN] = "1"; options_[Options::SNAPSHOT_NUM_RETAINED_MAX] = "1"; @@ -2306,7 +2617,7 @@ TEST_F(RealtimeWriteInteTest, TestReadMemoryBeforePrepareCommit) { RealtimeContextImpl::Cast(realtime_context)); ASSERT_NOK_WITH_MSG(realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket()), "ticket does not exist or has expired"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPlanExcludesRowsWrittenAfterMemoryEndOffset) { @@ -2350,7 +2661,7 @@ TEST_F(RealtimeWriteInteTest, TestPlanExcludesRowsWrittenAfterMemoryEndOffset) { ASSERT_OK_AND_ASSIGN(std::vector second_actual_rows, ReadRows(second_plan, realtime_context)); ASSERT_EQ(expected_rows, second_actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestReadFailsAfterRealtimeSplitTicketExpires) { @@ -2369,7 +2680,7 @@ TEST_F(RealtimeWriteInteTest, TestReadFailsAfterRealtimeSplitTicketExpires) { std::this_thread::sleep_for(std::chrono::milliseconds(50)); ASSERT_NOK_WITH_MSG(ReadRows(plan, realtime_context), "ticket does not exist or has expired"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestSuccessfulReaderCreationConsumesRealtimeSplitTicket) { @@ -2400,7 +2711,7 @@ TEST_F(RealtimeWriteInteTest, TestSuccessfulReaderCreationConsumesRealtimeSplitT ASSERT_NOK_WITH_MSG(table_read->CreateReader(plan->Splits()), "ticket does not exist or has expired"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTicket) { @@ -2432,10 +2743,10 @@ TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTick ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); ASSERT_EQ(rows, actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } -TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { +TEST_F(RealtimeWriteInteTest, TestCloseWriterWithBuildingDataInvalidatesContext) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); @@ -2445,10 +2756,28 @@ TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(rows, /*partitioned=*/false)); ASSERT_OK(writer->Write(std::move(batch))); - ASSERT_OK(writer->Close()); + ASSERT_NOK_WITH_MSG(writer->Close(), "not covered by a successful"); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); - ASSERT_EQ(rows, actual_rows); + ASSERT_NOK_WITH_MSG(ReadRows(realtime_context), "real-time context cannot be reused"); + ASSERT_NOK_WITH_MSG(CreateRealtimeWriter(realtime_context), + "real-time context cannot be reused"); +} + +TEST_F(RealtimeWriteInteTest, TestCloseWriterWithSealedDataInvalidatesContext) { + CreateTable(/*partition_keys=*/{}); + 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(MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"), + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK(writer->Seal()); + ASSERT_NOK_WITH_MSG(writer->Close(), "not covered by a successful"); + + ASSERT_NOK_WITH_MSG(CreateRealtimeWriter(realtime_context), + "real-time context cannot be reused"); } TEST_F(RealtimeWriteInteTest, TestPinnedPlanRemainsReadableAfterWriterClose) { @@ -2463,7 +2792,7 @@ TEST_F(RealtimeWriteInteTest, TestPinnedPlanRemainsReadableAfterWriterClose) { ASSERT_OK(writer->Write(std::move(batch))); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); ASSERT_EQ(rows, actual_rows); @@ -2531,7 +2860,7 @@ TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) { expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end()); ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestCountReaderDoesNotSupportRealtimeSplit) { @@ -2549,7 +2878,7 @@ TEST_F(RealtimeWriteInteTest, TestCountReaderDoesNotSupportRealtimeSplit) { CreatePlan(realtime_context, /*predicate=*/nullptr)); ASSERT_NOK_WITH_MSG(CountRows(plan, realtime_context), "does not support process-local real-time splits"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { @@ -2610,7 +2939,7 @@ TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { .ValueOrDie(); ASSERT_NE(nullptr, union_result); ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result)); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { @@ -2660,7 +2989,7 @@ TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { ASSERT_NE(nullptr, result); ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) << result->ToString(); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdown) { @@ -2714,7 +3043,7 @@ TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdown) { actual->ToString()); } } - PAIMON_RETURN_NOT_OK(writer->Close()); + PAIMON_RETURN_NOT_OK(PrepareAndClose(writer.get())); return Status::OK(); }; @@ -2798,7 +3127,7 @@ TEST_F(RealtimeWriteInteTest, TestMemoryBatchStatisticsPredicatePushdownWithDisk ASSERT_NE(nullptr, result); ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) << result->ToString(); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestNullPredicateForMemoryAndDisk) { @@ -2859,7 +3188,7 @@ TEST_F(RealtimeWriteInteTest, TestNullPredicateForMemoryAndDisk) { ASSERT_NE(nullptr, result); ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) << result->ToString(); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestUnionReadAfterColumnRename) { @@ -2913,7 +3242,7 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadAfterColumnRename) { ASSERT_NE(nullptr, result); ASSERT_TRUE(std::make_shared(expected)->Equals(*result)) << result->ToString(); - ASSERT_OK(second_writer->Close()); + ASSERT_OK(PrepareAndClose(second_writer.get())); } TEST_F(RealtimeWriteInteTest, TestUnionReadWithNestedStructProjection) { @@ -2963,7 +3292,7 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadWithNestedStructProjection) { [0, 2, [["beijing"]], "p0"], [0, 3, [["shenzhen"]], "p0"] ])"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestUnionReadWithSelectedMapKeys) { @@ -3037,7 +3366,7 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadWithSelectedMapKeys) { [0, 2, [40, null], "p0"], [0, 3, null, "p0"] ])"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); }; run_test(/*primary_key=*/false); ASSERT_FALSE(HasFatalFailure()); @@ -3088,8 +3417,12 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadWithVariantAccess) { arrow::ArrayVector disk_columns = {std::move(disk_offsets)}; disk_columns.insert(disk_columns.end(), disk_data->fields().begin(), disk_data->fields().end()); - const std::shared_ptr realtime_schema = - RealtimeOffsetUtils::CreateInputSchema(schema_); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create( + primary_key ? RealtimeStoreMode::PRIMARY_KEY : RealtimeStoreMode::APPEND_ONLY, + schema_)); + const std::shared_ptr& realtime_schema = schema_layout->InputSchema(); std::shared_ptr disk_realtime_data = arrow::StructArray::Make(std::move(disk_columns), realtime_schema->fields()) .ValueOrDie(); @@ -3151,7 +3484,7 @@ TEST_F(RealtimeWriteInteTest, TestUnionReadWithVariantAccess) { [0, 2, [30, "memory-a", null]], [0, 3, [40, "memory-b", "memory-fallback-b"]] ])"); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); }; run_test(/*primary_key=*/false); ASSERT_FALSE(HasFatalFailure()); @@ -3195,7 +3528,7 @@ TEST_F(RealtimeWriteInteTest, TestRefreshCommittedSnapshotReclaimsMemory) { ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh, GetRealtimeMemoryUsage(realtime_context)); ASSERT_LT(memory_usage_after_refresh, memory_usage_before_refresh); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestPlanPinsMemoryAcrossRefresh) { @@ -3233,7 +3566,7 @@ TEST_F(RealtimeWriteInteTest, TestPlanPinsMemoryAcrossRefresh) { ASSERT_OK_AND_ASSIGN(std::vector refreshed_rows, ReadRows(refreshed_plan, realtime_context)); ASSERT_EQ(pinned_rows, refreshed_rows); - ASSERT_OK(writer->Close()); + ASSERT_OK(PrepareAndClose(writer.get())); } TEST_F(RealtimeWriteInteTest, TestReaderPinsMemoryAcrossRefresh) { @@ -3418,7 +3751,7 @@ TEST_F(RealtimeWriteInteTest, TestOverwriteRequiresReopenRealtimeContext) { ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_failed_refresh, GetRealtimeMemoryUsage(realtime_context)); ASSERT_EQ(memory_usage_before_overwrite, memory_usage_after_failed_refresh); - ASSERT_OK(writer->Close()); + ASSERT_NOK_WITH_MSG(writer->Close(), "not covered by a successful"); writer.reset(); realtime_context.reset(); @@ -3900,7 +4233,7 @@ TEST_F(RealtimeWriteInteTest, TestDropPartitionRequiresReopenRealtimeContext) { ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh, GetRealtimeMemoryUsage(realtime_context)); ASSERT_EQ(memory_usage_before_refresh, memory_usage_after_refresh); - ASSERT_OK(writer->Close()); + ASSERT_NOK_WITH_MSG(writer->Close(), "not covered by a successful"); writer.reset(); realtime_context.reset();