From b935705a8db08cdf345ef1fb3bbb6b762e94b0e9 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 1 Aug 2026 12:33:30 +0800 Subject: [PATCH 1/4] feat: support real-time append writes with pluggable memory indexers --- include/paimon/api.h | 5 + include/paimon/file_store_commit.h | 16 + include/paimon/file_store_write.h | 16 + .../realtime/arrow_mem_indexer_factory.h | 32 ++ include/paimon/realtime/mem_indexer.h | 116 +++++ .../realtime/realtime_commit_progress.h | 45 ++ include/paimon/realtime/realtime_context.h | 67 +++ include/paimon/type_fwd.h | 1 + include/paimon/utils/special_field_ids.h | 3 + include/paimon/write_context.h | 16 + src/paimon/CMakeLists.txt | 6 + src/paimon/common/table/special_fields.h | 8 +- .../common/table/special_fields_test.cpp | 7 + .../append/append_compact_coordinator.cpp | 5 +- .../core/memory/writer_memory_manager.h | 32 +- .../memory/writer_memory_manager_test.cpp | 15 + .../operation/abstract_file_store_write.cpp | 77 +++- .../operation/abstract_file_store_write.h | 10 + .../core/operation/abstract_split_read.cpp | 8 +- .../core/operation/abstract_split_read.h | 7 +- .../operation/abstract_split_read_test.cpp | 50 +-- .../append_only_file_store_write.cpp | 46 +- .../operation/append_only_file_store_write.h | 10 + .../append_only_file_store_write_test.cpp | 117 +++++ .../commit/realtime_snapshot_properties.cpp | 287 ++++++++++++ .../commit/realtime_snapshot_properties.h | 75 ++++ .../realtime_snapshot_properties_test.cpp | 363 +++++++++++++++ .../operation/data_evolution_split_read.cpp | 5 +- .../core/operation/file_store_commit_impl.cpp | 71 ++- .../core/operation/file_store_commit_impl.h | 5 +- .../operation/file_store_commit_impl_test.cpp | 16 +- .../core/operation/file_store_write.cpp | 31 +- .../core/operation/internal_read_context.cpp | 6 + src/paimon/core/operation/write_context.cpp | 12 +- .../core/realtime/arrow_mem_indexer.cpp | 204 +++++++++ src/paimon/core/realtime/arrow_mem_indexer.h | 68 +++ .../realtime/arrow_mem_indexer_factory.cpp | 42 ++ src/paimon/core/realtime/partition_bucket.h | 52 +++ .../realtime/realtime_append_only_writer.cpp | 231 ++++++++++ .../realtime/realtime_append_only_writer.h | 90 ++++ src/paimon/core/realtime/realtime_context.cpp | 89 ++++ src/paimon/core/utils/commit_increment.h | 11 + test/inte/CMakeLists.txt | 7 + test/inte/realtime_write_inte_test.cpp | 423 ++++++++++++++++++ 44 files changed, 2731 insertions(+), 72 deletions(-) create mode 100644 include/paimon/realtime/arrow_mem_indexer_factory.h create mode 100644 include/paimon/realtime/mem_indexer.h create mode 100644 include/paimon/realtime/realtime_commit_progress.h create mode 100644 include/paimon/realtime/realtime_context.h create mode 100644 src/paimon/core/operation/commit/realtime_snapshot_properties.cpp create mode 100644 src/paimon/core/operation/commit/realtime_snapshot_properties.h create mode 100644 src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp create mode 100644 src/paimon/core/realtime/arrow_mem_indexer.cpp create mode 100644 src/paimon/core/realtime/arrow_mem_indexer.h create mode 100644 src/paimon/core/realtime/arrow_mem_indexer_factory.cpp create mode 100644 src/paimon/core/realtime/partition_bucket.h create mode 100644 src/paimon/core/realtime/realtime_append_only_writer.cpp create mode 100644 src/paimon/core/realtime/realtime_append_only_writer.h create mode 100644 src/paimon/core/realtime/realtime_context.cpp create mode 100644 test/inte/realtime_write_inte_test.cpp diff --git a/include/paimon/api.h b/include/paimon/api.h index 1771a4d9..a0843b3e 100644 --- a/include/paimon/api.h +++ b/include/paimon/api.h @@ -38,5 +38,10 @@ #include "paimon/table/source/table_scan.h" // IWYU pragma: export #include "paimon/write_context.h" // IWYU pragma: export +// IWYU pragma: begin_exports +#include "paimon/realtime/mem_indexer.h" +#include "paimon/realtime/realtime_context.h" +// IWYU pragma: end_exports + /// Top-level namespace for Paimon C++ API. namespace paimon {} diff --git a/include/paimon/file_store_commit.h b/include/paimon/file_store_commit.h index 82bbc003..7ef09c5c 100644 --- a/include/paimon/file_store_commit.h +++ b/include/paimon/file_store_commit.h @@ -30,6 +30,7 @@ #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" #include "paimon/metrics.h" +#include "paimon/realtime/realtime_commit_progress.h" #include "paimon/result.h" #include "paimon/status.h" #include "paimon/type_fwd.h" @@ -71,6 +72,21 @@ class PAIMON_EXPORT FileStoreCommit { int64_t commit_identifier = BATCH_WRITE_COMMIT_IDENTIFIER, std::optional watermark = std::nullopt) = 0; + /// Commit sealed real-time segments and persist their partition-bucket offset progress. + /// + /// Entries for each partition-bucket must form a contiguous range beginning after the offset + /// recorded by the latest committed snapshot. Input entries may be unordered; this method + /// orders them by partition, bucket, and offset before validating continuity. The resulting + /// snapshot atomically publishes the data files and the updated offset map. + /// + /// @param realtime_commits Commit messages and inclusive offset ranges to commit. + /// @param commit_identifier Identifier of the streaming commit operation. + /// @param watermark Optional event-time watermark. + /// @return Status indicating the success or failure of the commit operation. + virtual Status CommitWithProgress(const std::vector& realtime_commits, + int64_t commit_identifier, + std::optional watermark) = 0; + /// Filter out all `std::vector` which have been committed and commit the /// remaining ones. /// diff --git a/include/paimon/file_store_write.h b/include/paimon/file_store_write.h index e870c442..d79b2595 100644 --- a/include/paimon/file_store_write.h +++ b/include/paimon/file_store_write.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "paimon/commit_message.h" @@ -28,6 +29,7 @@ #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" #include "paimon/metrics.h" +#include "paimon/realtime/realtime_commit_progress.h" #include "paimon/result.h" #include "paimon/status.h" #include "paimon/type_fwd.h" @@ -85,8 +87,22 @@ class PAIMON_EXPORT FileStoreWrite { /// /// @return A Result containing `std::vector>` objects, /// representing the generated commit messages. + /// @note Real-time writers must use `PrepareCommitWithProgress()` so offset ranges are not + /// discarded. virtual Result>> PrepareCommit( bool wait_compaction = true, int64_t commit_identifier = BATCH_WRITE_COMMIT_IDENTIFIER) = 0; + + /// 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. + /// + /// @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. + virtual Result> PrepareCommitWithProgress( + int64_t commit_identifier); + virtual std::shared_ptr GetMetrics() const = 0; virtual Status Close() = 0; }; diff --git a/include/paimon/realtime/arrow_mem_indexer_factory.h b/include/paimon/realtime/arrow_mem_indexer_factory.h new file mode 100644 index 00000000..efa11767 --- /dev/null +++ b/include/paimon/realtime/arrow_mem_indexer_factory.h @@ -0,0 +1,32 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "paimon/realtime/mem_indexer.h" + +namespace paimon { + +/// Factory for Paimon's default Arrow-backed `MemIndexer`. +class PAIMON_EXPORT ArrowMemIndexerFactory : public MemIndexerFactory { + public: + /// Creates an Arrow-backed indexer for one partition and bucket. + Result> Create( + ::ArrowSchema* write_schema, const std::map& options, + const std::shared_ptr& memory_pool) override; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/mem_indexer.h b/include/paimon/realtime/mem_indexer.h new file mode 100644 index 00000000..18cd13ea --- /dev/null +++ b/include/paimon/realtime/mem_indexer.h @@ -0,0 +1,116 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/reader/batch_reader.h" +#include "paimon/record_batch.h" +#include "paimon/result.h" +#include "paimon/utils/range.h" +#include "paimon/visibility.h" + +struct ArrowSchema; + +namespace paimon { + +class MemoryPool; + +/// A record batch and the contiguous offset range assigned to its rows. +/// +/// The batch contains only the table write fields. `_OFFSET` is carried separately by +/// `offset_range`; row `i` corresponds to `offset_range.from + i`. Paimon adds the physical +/// `_OFFSET` column when the sealed segment is written to data files. +struct PAIMON_EXPORT RealtimeWriteBatch { + /// Input batch whose ownership is transferred to `MemIndexer::Write`. + std::unique_ptr batch; + /// Inclusive `[from, to]` offset range covered by `batch`. + Range offset_range; +}; + +/// Opaque handle to an immutable segment returned by `MemIndexer::SealForCommit`. +/// +/// A plugin may store the segment in memory or in spill files. Callers use this handle only to +/// request commit readers and inspect its offset range. +class PAIMON_EXPORT RealtimeSegmentHandle { + public: + virtual ~RealtimeSegmentHandle() = default; + + /// Returns the inclusive offset range covered by this segment. + virtual Range GetOffsetRange() const = 0; +}; + +/// Plugin interface for buffering real-time writes before Paimon data-file generation. +/// +/// Paimon serializes calls to `Write` and `SealForCommit` for the same indexer. After sealing, +/// `CreateCommitReaders` may read the immutable sealed segment while later `Write` calls append to +/// a new building segment. Paimon retains control of file format, rolling, indexes, and +/// commit-message generation. +class PAIMON_EXPORT MemIndexer { + public: + virtual ~MemIndexer() = default; + + /// Adds a batch to the current building segment. + /// + /// The number of rows must equal the size of `offset_range`. + virtual Status Write(RealtimeWriteBatch&& batch) = 0; + + /// Seals the current building data and opens a new building segment. + /// + /// Returns an immutable segment handle, or `std::nullopt` when there is no data to seal. + virtual Result>> SealForCommit() = 0; + + /// Creates readers that expose all rows in a sealed segment for Paimon file writing. + /// + /// Concatenating the returned readers must produce every sealed row exactly once and in write + /// order. Each output batch contains `_VALUE_KIND` followed by all fields from the factory's + /// `write_schema`; it does not contain `_OFFSET`. + virtual Result>> CreateCommitReaders( + const std::shared_ptr& segment) = 0; + + /// Returns the number of bytes currently retained by the building segment. + virtual uint64_t GetMemoryUsage() const = 0; + + /// Releases resources owned by this indexer and rejects subsequent writes or seals. + virtual Status Close() = 0; +}; + +/// Factory for application-provided `MemIndexer` implementations. +class PAIMON_EXPORT MemIndexerFactory { + public: + virtual ~MemIndexerFactory() = default; + + /// Creates an indexer configured with the supplied schema, options, and memory pool. + /// + /// `write_schema` uses the Arrow C Data Interface and contains the table fields accepted by + /// `Write`. It is valid only during this call. An implementation may consume its contents by + /// using an Arrow C Data Interface importer; otherwise Paimon releases them after this method + /// returns. + /// @param write_schema Table write schema without Paimon-generated real-time fields. + /// @param options Effective table options available to the indexer. + /// @param memory_pool Memory pool provided by the write context. + virtual Result> Create( + ::ArrowSchema* write_schema, const std::map& options, + const std::shared_ptr& memory_pool) = 0; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/realtime_commit_progress.h b/include/paimon/realtime/realtime_commit_progress.h new file mode 100644 index 00000000..6f346ca8 --- /dev/null +++ b/include/paimon/realtime/realtime_commit_progress.h @@ -0,0 +1,45 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/commit_message.h" +#include "paimon/utils/range.h" +#include "paimon/visibility.h" + +namespace paimon { + +/// A real-time commit message and its partition-bucket offset progress. +/// +/// Offsets are scoped to one partition and bucket. `offset_range` is inclusive and covers all +/// rows represented by `commit_message`. The progress fields are not embedded in +/// `CommitMessage` serialization. +struct PAIMON_EXPORT RealtimeCommitProgress { + /// Paimon commit message generated from one sealed segment. + std::shared_ptr commit_message; + /// Partition path such as `dt=2`, or an empty string for an unpartitioned table. + std::string partition; + /// Bucket containing the sealed segment. + int32_t bucket; + /// Inclusive offset range represented by the commit message. + Range offset_range; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h new file mode 100644 index 00000000..8f4db839 --- /dev/null +++ b/include/paimon/realtime/realtime_context.h @@ -0,0 +1,67 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/result.h" +#include "paimon/visibility.h" + +struct ArrowSchema; + +namespace paimon { + +class MemIndexer; +class MemIndexerFactory; +class MemoryPool; + +/// Shared context that owns the `MemIndexer` instances used by a real-time writer. +/// +/// Applications attach one context to `WriteContext`. The context uses either the default Arrow +/// implementation or an application-provided factory and keeps each created indexer available +/// across multiple writes and prepare-commit operations. +class PAIMON_EXPORT RealtimeContext { + public: + /// Creates a context backed by Paimon's default Arrow `MemIndexer`. + static Result> Create(); + + /// Creates a context backed by an application-provided indexer factory. + /// + /// @param factory Non-null factory used to create indexers on demand. + static Result> Create( + const std::shared_ptr& factory); + + /// Returns the stable indexer associated with a partition and bucket, creating it if needed. + Result> GetOrCreateMemIndexer( + const std::string& partition, int32_t bucket, std::unique_ptr<::ArrowSchema> write_schema, + const std::map& options, + const std::shared_ptr& memory_pool); + + ~RealtimeContext(); + + private: + class Impl; + + explicit RealtimeContext(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/include/paimon/type_fwd.h b/include/paimon/type_fwd.h index d9429f16..6656ff0c 100644 --- a/include/paimon/type_fwd.h +++ b/include/paimon/type_fwd.h @@ -56,6 +56,7 @@ class FileStatus; class BasicFileStatus; class RecordBatch; +class RealtimeContext; class FormatWriter; diff --git a/include/paimon/utils/special_field_ids.h b/include/paimon/utils/special_field_ids.h index c4e03f00..48abcc85 100644 --- a/include/paimon/utils/special_field_ids.h +++ b/include/paimon/utils/special_field_ids.h @@ -42,6 +42,9 @@ class SpecialFieldIds { /// Special field ID reserved for index score. Value: CPP_FIELD_ID_END - 1 inline static constexpr int32_t INDEX_SCORE = CPP_FIELD_ID_END - 1; + // TODO(xinyu.lxy): Add _OFFSET to the Java codebase so compaction preserves its information. + /// Special field ID reserved for real-time offset. Value: CPP_FIELD_ID_END - 2 + inline static constexpr int32_t OFFSET = CPP_FIELD_ID_END - 2; }; } // namespace paimon diff --git a/include/paimon/write_context.h b/include/paimon/write_context.h index 675befbd..c99dabf7 100644 --- a/include/paimon/write_context.h +++ b/include/paimon/write_context.h @@ -40,6 +40,9 @@ class MemoryPool; /// @see WriteContextBuilder class PAIMON_EXPORT WriteContext { public: + /// Constructs a write context with an optional real-time write context. + /// + /// Prefer `WriteContextBuilder`. WriteContext(const std::string& root_path, const std::string& commit_user, bool is_streaming_mode, bool ignore_num_bucket_check, bool ignore_previous_files, bool enable_multi_thread_spill, const std::optional& write_id, @@ -48,6 +51,7 @@ class PAIMON_EXPORT WriteContext { const std::shared_ptr& executor, const std::string& temp_directory, const std::shared_ptr& specific_file_system, const std::map& fs_scheme_to_identifier_map, + const std::shared_ptr& realtime_context, const std::map& options); ~WriteContext(); @@ -112,6 +116,11 @@ class PAIMON_EXPORT WriteContext { return enable_multi_thread_spill_; } + /// Returns the configured real-time context, or `nullptr` when real-time writing is disabled. + std::shared_ptr GetRealtimeContext() const { + return realtime_context_; + } + private: std::string root_path_; std::string commit_user_; @@ -127,6 +136,7 @@ class PAIMON_EXPORT WriteContext { std::string temp_directory_; std::shared_ptr specific_file_system_; std::map fs_scheme_to_identifier_map_; + std::shared_ptr realtime_context_; std::map options_; }; @@ -210,6 +220,12 @@ class PAIMON_EXPORT WriteContextBuilder { /// @note If not set, use default file system (configured in `Options::FILE_SYSTEM`) WriteContextBuilder& WithFileSystem(const std::shared_ptr& file_system); + /// Enables the real-time write path with the provided shared context. + /// @param realtime_context Non-null context that owns the real-time indexers. + /// @return Reference to this builder for method chaining. + WriteContextBuilder& WithRealtimeContext( + const std::shared_ptr& realtime_context); + /// Sets a mapping from URI schemes (e.g., "file", "oss") to registered file system /// identifiers. This allows selecting different pre-registered file system implementations /// based on the URI scheme at runtime. diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index ce2f5ddf..b236676f 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -324,6 +324,7 @@ set(PAIMON_CORE_SRCS core/operation/commit/commit_changes_provider.cpp core/operation/commit/compacted_changelog_path_resolver.cpp core/operation/commit/overwrite_changes_provider.cpp + core/operation/commit/realtime_snapshot_properties.cpp core/operation/commit/row_id_column_conflict_checker.cpp core/operation/commit/manifest_entry_changes.cpp core/operation/commit/row_tracking_commit_utils.cpp @@ -349,6 +350,10 @@ set(PAIMON_CORE_SRCS core/manifest/snapshot_live_manifest_entries.cpp core/operation/write_context.cpp core/operation/write_restore.cpp + core/realtime/arrow_mem_indexer.cpp + core/realtime/arrow_mem_indexer_factory.cpp + core/realtime/realtime_append_only_writer.cpp + core/realtime/realtime_context.cpp core/postpone/postpone_bucket_writer.cpp core/schema/arrow_schema_validator.cpp core/schema/schema_manager.cpp @@ -763,6 +768,7 @@ if(PAIMON_BUILD_TESTS) core/operation/commit/overwrite_changes_provider_test.cpp core/operation/commit/row_id_column_conflict_checker_test.cpp core/operation/commit/row_tracking_commit_utils_test.cpp + core/operation/commit/realtime_snapshot_properties_test.cpp core/operation/commit/sequence_snapshot_properties_test.cpp core/operation/commit/retry_waiter_test.cpp core/operation/key_value_file_store_write_test.cpp diff --git a/src/paimon/common/table/special_fields.h b/src/paimon/common/table/special_fields.h index 74b95b19..612c43f1 100644 --- a/src/paimon/common/table/special_fields.h +++ b/src/paimon/common/table/special_fields.h @@ -60,6 +60,12 @@ struct SpecialFields { return data_field; } + static const DataField& Offset() { + static const DataField data_field = + DataField(SpecialFieldIds::OFFSET, arrow::field("_OFFSET", arrow::int64())); + return data_field; + } + static const DataField& IndexScore() { static const DataField data_field = DataField(SpecialFieldIds::INDEX_SCORE, arrow::field("_INDEX_SCORE", arrow::float32())); @@ -72,7 +78,7 @@ struct SpecialFields { } return field_name == SequenceNumber().Name() || field_name == ValueKind().Name() || field_name == RowKind().Name() || field_name == RowId().Name() || - field_name == IndexScore().Name(); + field_name == Offset().Name() || field_name == IndexScore().Name(); } // TODO(xinyu.lxy): add a func to complete row-tracking fields diff --git a/src/paimon/common/table/special_fields_test.cpp b/src/paimon/common/table/special_fields_test.cpp index 68e805fd..3dd8b670 100644 --- a/src/paimon/common/table/special_fields_test.cpp +++ b/src/paimon/common/table/special_fields_test.cpp @@ -49,6 +49,12 @@ TEST(SpecialFieldsTest, TestRowIdField) { ASSERT_EQ(SpecialFields::RowId().Type()->id(), arrow::Type::INT64); } +TEST(SpecialFieldsTest, TestOffsetField) { + ASSERT_EQ(SpecialFields::Offset().Id(), std::numeric_limits::max() - 10000 - 2); + ASSERT_EQ(SpecialFields::Offset().Name(), "_OFFSET"); + ASSERT_EQ(SpecialFields::Offset().Type()->id(), arrow::Type::INT64); +} + TEST(SpecialFieldsTest, TestIndexScore) { ASSERT_EQ(SpecialFields::IndexScore().Id(), std::numeric_limits::max() - 10000 - 1); ASSERT_EQ(SpecialFields::IndexScore().Name(), "_INDEX_SCORE"); @@ -65,6 +71,7 @@ TEST(SpecialFieldsTest, TestIsSystemField) { ASSERT_FALSE(SpecialFields::IsSystemField("VALUE_KIND")); ASSERT_TRUE(SpecialFields::IsSystemField("rowkind")); ASSERT_TRUE(SpecialFields::IsSystemField("_ROW_ID")); + ASSERT_TRUE(SpecialFields::IsSystemField("_OFFSET")); ASSERT_TRUE(SpecialFields::IsSystemField("_INDEX_SCORE")); ASSERT_TRUE(SpecialFields::IsSystemField("_KEY_0")); } diff --git a/src/paimon/core/append/append_compact_coordinator.cpp b/src/paimon/core/append/append_compact_coordinator.cpp index 37693a58..96d13d1d 100644 --- a/src/paimon/core/append/append_compact_coordinator.cpp +++ b/src/paimon/core/append/append_compact_coordinator.cpp @@ -19,7 +19,6 @@ #include "paimon/append/append_compact_coordinator.h" #include -#include #include #include #include @@ -162,6 +161,7 @@ std::unique_ptr CreateFileStoreWrite( const std::shared_ptr& arrow_schema, const std::shared_ptr& partition_schema, const CoreOptions& core_options, const std::shared_ptr& executor, const std::shared_ptr& pool) { + const RealtimeSnapshotProperties::OffsetMap realtime_committed_offsets; return std::make_unique( path_factory, snapshot_manager, schema_manager, /*commit_user=*/"compact-coordinator", @@ -171,7 +171,8 @@ std::unique_ptr CreateFileStoreWrite( /*io_manager=*/nullptr, core_options, /*ignore_previous_files=*/true, /*is_streaming_mode=*/false, - /*ignore_num_bucket_check=*/false, executor, pool); + /*ignore_num_bucket_check=*/false, + /*realtime_context=*/nullptr, realtime_committed_offsets, executor, pool); } /// Load schema from table path and merge user options with schema options. diff --git a/src/paimon/core/memory/writer_memory_manager.h b/src/paimon/core/memory/writer_memory_manager.h index 811667cc..093c4ac4 100644 --- a/src/paimon/core/memory/writer_memory_manager.h +++ b/src/paimon/core/memory/writer_memory_manager.h @@ -33,16 +33,17 @@ class BatchWriter; class WriterMemoryManager { public: explicit WriterMemoryManager(uint64_t memory_limit) : memory_limit_(memory_limit) {} + virtual ~WriterMemoryManager() = default; /// Register a writer when create a new `BatchWriter` in `AbstractFileStoreWrite::GetWriter()` - void RegisterWriter(BatchWriter* writer); + virtual void RegisterWriter(BatchWriter* writer); /// Unregister a writer when the `BatchWriter` has been erased in `AbstractFileStoreWrite` - void UnregisterWriter(BatchWriter* writer); + virtual void UnregisterWriter(BatchWriter* writer); /// Refresh the memory usage of a writer when after `BatchWriter::PrepareCommit()` - void RefreshWriterMemory(BatchWriter* writer); + virtual void RefreshWriterMemory(BatchWriter* writer); /// Check if the total memory usage exceeds the limit after `BatchWriter::Write()`, and trigger /// flush if needed. - Status OnWriteCompleted(BatchWriter* writer); + virtual Status OnWriteCompleted(BatchWriter* writer); private: struct Candidate { @@ -59,4 +60,27 @@ class WriterMemoryManager { std::unordered_map writer_memory_; }; +/// Memory manager policy for writers whose memory is managed outside `AbstractFileStoreWrite`. +class NoopWriterMemoryManager final : public WriterMemoryManager { + public: + NoopWriterMemoryManager() : WriterMemoryManager(/*memory_limit=*/0) {} + + void RegisterWriter(BatchWriter* writer) override { + (void)writer; + } + + void UnregisterWriter(BatchWriter* writer) override { + (void)writer; + } + + void RefreshWriterMemory(BatchWriter* writer) override { + (void)writer; + } + + Status OnWriteCompleted(BatchWriter* writer) override { + (void)writer; + return Status::OK(); + } +}; + } // namespace paimon diff --git a/src/paimon/core/memory/writer_memory_manager_test.cpp b/src/paimon/core/memory/writer_memory_manager_test.cpp index 1a91bf64..9f5eb469 100644 --- a/src/paimon/core/memory/writer_memory_manager_test.cpp +++ b/src/paimon/core/memory/writer_memory_manager_test.cpp @@ -245,4 +245,19 @@ TEST(WriterMemoryManagerTest, ReturnsConfigurationErrorWhenNoWriterCanReleaseEno ASSERT_EQ(flush_history, std::vector({"writer_a"})); } +TEST(WriterMemoryManagerTest, NoopManagerDoesNotTrackOrFlushWriter) { + NoopWriterMemoryManager manager; + std::vector flush_history; + FakeBatchWriter writer("writer", &flush_history); + writer.SetMemoryUsage(100); + + manager.RegisterWriter(&writer); + ASSERT_OK(manager.OnWriteCompleted(&writer)); + manager.RefreshWriterMemory(&writer); + manager.UnregisterWriter(&writer); + + ASSERT_TRUE(flush_history.empty()); + ASSERT_EQ(writer.GetMemoryUsage(), 100); +} + } // namespace paimon::test diff --git a/src/paimon/core/operation/abstract_file_store_write.cpp b/src/paimon/core/operation/abstract_file_store_write.cpp index d7dafa18..14ff04be 100644 --- a/src/paimon/core/operation/abstract_file_store_write.cpp +++ b/src/paimon/core/operation/abstract_file_store_write.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include "fmt/format.h" #include "paimon/common/data/binary_row.h" @@ -30,6 +31,7 @@ #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/partition_bucket.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" #include "paimon/core/table/bucket_mode.h" @@ -133,8 +135,7 @@ Status AbstractFileStoreWrite::Write(std::unique_ptr&& batch) { GetWriter(partition, batch->GetBucket())); assert(writer); PAIMON_RETURN_NOT_OK(writer->Write(std::move(batch))); - PAIMON_RETURN_NOT_OK(writer_memory_manager_->OnWriteCompleted(writer.get())); - return Status::OK(); + return writer_memory_manager_->OnWriteCompleted(writer.get()); } Status AbstractFileStoreWrite::Compact(const std::map& partition, @@ -147,6 +148,10 @@ Status AbstractFileStoreWrite::Compact(const std::map& Result>> AbstractFileStoreWrite::PrepareCommit( bool wait_compaction, int64_t commit_identifier) { + if (IsRealtimeWrite()) { + return Status::Invalid( + "real-time writer must use PrepareCommitWithProgress to preserve offset ranges"); + } if (batch_committed_) { return Status::Invalid("batch write mode only support one-time committing."); } @@ -261,7 +266,74 @@ Result>> AbstractFileStoreWrite::Prep return result; } +Result> AbstractFileStoreWrite::PrepareCommitWithProgress( + int64_t) { + if (!IsRealtimeWrite()) { + return Status::Invalid("PrepareCommitWithProgress is only supported by a real-time writer"); + } + if (is_streaming_mode_ == false) { + return Status::Invalid("PrepareCommitWithProgress requires streaming mode"); + } + // 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> AbstractFileStoreWrite::PrepareRealtimeCommit() { + struct WriterSnapshot { + BinaryRow partition; + int32_t bucket; + int32_t total_buckets; + std::shared_ptr writer; + }; + std::vector writer_snapshots; + { + std::lock_guard lock(writers_mutex_); + for (const auto& [partition, buckets] : writers_) { + for (const auto& [bucket, container] : buckets) { + writer_snapshots.push_back( + WriterSnapshot{partition, bucket, container.total_buckets, container.writer}); + } + } + } + + std::vector result; + for (const WriterSnapshot& snapshot : writer_snapshots) { + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, + snapshot.writer->PrepareCommit(/*wait_compaction=*/false)); + writer_memory_manager_->RefreshWriterMemory(snapshot.writer.get()); + auto committable = std::make_shared( + snapshot.partition, snapshot.bucket, snapshot.total_buckets, + increment.GetNewFilesIncrement(), increment.GetCompactIncrement()); + if (!increment.GetRealtimeOffsetRange()) { + if (!committable->IsEmpty()) { + return Status::Invalid("real-time commit message does not have an offset range"); + } + continue; + } + if (committable->IsEmpty()) { + return Status::Invalid("sealed real-time segment produced an empty commit message"); + } + PAIMON_ASSIGN_OR_RAISE(std::string partition, + file_store_path_factory_->GetPartitionString(snapshot.partition)); + partition = PartitionBucket::NormalizePartition(std::move(partition)); + result.push_back(RealtimeCommitProgress{committable, std::move(partition), snapshot.bucket, + increment.GetRealtimeOffsetRange().value()}); + } + { + std::lock_guard lock(realtime_metrics_mutex_); + std::shared_ptr metrics = compaction_metrics_->GetMetrics(); + for (const WriterSnapshot& snapshot : writer_snapshots) { + metrics->Merge(snapshot.writer->GetMetrics()); + } + metrics_->Overwrite(metrics); + } + return result; +} + Status AbstractFileStoreWrite::Close() { + std::lock_guard lock(writers_mutex_); for (auto& [_, bucket_writers] : writers_) { for (auto& [_, writer_container] : bucket_writers) { writer_memory_manager_->UnregisterWriter(writer_container.writer.get()); @@ -325,6 +397,7 @@ Result> AbstractFileStoreWrite::ScanExistingFileMe Result> AbstractFileStoreWrite::GetWriter(const BinaryRow& partition, int32_t bucket) { + std::lock_guard lock(writers_mutex_); auto partition_iter = writers_.find(partition); if (partition_iter != writers_.end()) { auto& buckets = partition_iter->second; diff --git a/src/paimon/core/operation/abstract_file_store_write.h b/src/paimon/core/operation/abstract_file_store_write.h index ca6016d2..f43046bf 100644 --- a/src/paimon/core/operation/abstract_file_store_write.h +++ b/src/paimon/core/operation/abstract_file_store_write.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -88,6 +89,8 @@ class AbstractFileStoreWrite : public FileStoreWrite { Result>> PrepareCommit( bool wait_compaction, int64_t commit_identifier) override; + Result> PrepareCommitWithProgress( + int64_t commit_identifier) override; Status Close() override; std::shared_ptr GetMetrics() const override; @@ -116,6 +119,10 @@ class AbstractFileStoreWrite : public FileStoreWrite { virtual Result> CreateFileStoreScan( const std::shared_ptr& filter) const = 0; + virtual bool IsRealtimeWrite() const { + return false; + } + Result> ScanExistingFileMetas(const BinaryRow& partition, int32_t bucket) const; int32_t GetDefaultBucketNum() const; @@ -142,10 +149,13 @@ class AbstractFileStoreWrite : public FileStoreWrite { private: Result> GetWriter(const BinaryRow& partition, int32_t bucket); + Result> PrepareRealtimeCommit(); private: std::unordered_map>> writers_; + std::mutex writers_mutex_; + std::mutex realtime_metrics_mutex_; bool ignore_previous_files_ = false; bool is_streaming_mode_ = false; bool ignore_num_bucket_check_ = false; diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 0850d943..624999be 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -191,9 +191,8 @@ Result> AbstractSplitRead::CreateFieldMappingRe PAIMON_ASSIGN_OR_RAISE(field_mapping, field_mapping_builder->CreateFieldMapping(file_fields)); } else { - PAIMON_ASSIGN_OR_RAISE( - std::vector projected_data_fields, - ProjectFieldsForRowTrackingAndDataEvolution(data_schema, file_meta->write_cols)); + PAIMON_ASSIGN_OR_RAISE(std::vector projected_data_fields, + BuildDataFieldsForFieldMapping(data_schema, file_meta->write_cols)); auto converted_fields = BlobUtils::ConvertBlobInlineDataFields(projected_data_fields, blob_inline_fields); PAIMON_ASSIGN_OR_RAISE(field_mapping, @@ -329,13 +328,14 @@ Result> AbstractSplitRead::ApplyVariantShreddin return std::move(file_reader); } -Result> AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( +Result> AbstractSplitRead::BuildDataFieldsForFieldMapping( const std::shared_ptr& data_schema, const std::optional>& write_cols) { std::vector projected_fields; const std::vector& partition_keys = data_schema->PartitionKeys(); if (write_cols == std::nullopt) { projected_fields = data_schema->Fields(); + projected_fields.insert(projected_fields.begin(), SpecialFields::Offset()); } else { if (write_cols.value().empty()) { return Status::Invalid("write cols cannot be empty"); diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index ea3f9070..5385f942 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -93,10 +93,9 @@ class AbstractSplitRead : public SplitRead { const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const = 0; - // 1. project write cols to data schema - // 2. add partition fields (if write cols not contain) - // 3. add row tracking fields - static Result> ProjectFieldsForRowTrackingAndDataEvolution( + // Builds the fields used to create a file field mapping from write columns, partition columns, + // and internal fields. + static Result> BuildDataFieldsForFieldMapping( const std::shared_ptr& data_schema, const std::optional>& write_cols); diff --git a/src/paimon/core/operation/abstract_split_read_test.cpp b/src/paimon/core/operation/abstract_split_read_test.cpp index 6fdd7f26..65020832 100644 --- a/src/paimon/core/operation/abstract_split_read_test.cpp +++ b/src/paimon/core/operation/abstract_split_read_test.cpp @@ -48,7 +48,7 @@ TEST(AbstractSplitReadTest, TestNeedCompleteRowTrackingFields) { arrow::schema(fields))); } -TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { +TEST(AbstractSplitReadTest, TestBuildDataFieldsForFieldMapping) { { // test no partition std::vector fields = {DataField(0, arrow::field("name", arrow::utf8())), @@ -62,10 +62,10 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test write_cols is std::nullopt - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, /*write_cols=*/std::nullopt)); - std::vector expected = fields; + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, /*write_cols=*/std::nullopt)); + std::vector expected = {SpecialFields::Offset()}; + expected.insert(expected.end(), fields.begin(), fields.end()); expected.push_back(SpecialFields::RowId()); expected.push_back(SpecialFields::SequenceNumber()); ASSERT_EQ(result, expected); @@ -73,9 +73,8 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test with write_cols std::vector write_cols = {"name", "age"}; - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols)); + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, write_cols)); std::vector expected = {fields[0], fields[2], SpecialFields::RowId(), SpecialFields::SequenceNumber()}; ASSERT_EQ(result, expected); @@ -83,9 +82,9 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test with empty write_cols std::vector write_cols = {}; - ASSERT_NOK_WITH_MSG(AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols), - "write cols cannot be empty"); + ASSERT_NOK_WITH_MSG( + AbstractSplitRead::BuildDataFieldsForFieldMapping(table_schema, write_cols), + "write cols cannot be empty"); } } { @@ -102,10 +101,10 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test write_cols is std::nullopt - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, /*write_cols=*/std::nullopt)); - std::vector expected = fields; + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, /*write_cols=*/std::nullopt)); + std::vector expected = {SpecialFields::Offset()}; + expected.insert(expected.end(), fields.begin(), fields.end()); expected.push_back(SpecialFields::RowId()); expected.push_back(SpecialFields::SequenceNumber()); ASSERT_EQ(result, expected); @@ -113,9 +112,8 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test with write_cols and write_cols not contain partition fields std::vector write_cols = {"name", "age"}; - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols)); + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, write_cols)); std::vector expected = {fields[0], fields[3], fields[1], SpecialFields::RowId(), SpecialFields::SequenceNumber()}; @@ -124,9 +122,8 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test with write_cols and write_cols contain partition fields std::vector write_cols = {"age", "name", "ds"}; - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols)); + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, write_cols)); std::vector expected = {fields[3], fields[0], fields[1], SpecialFields::RowId(), SpecialFields::SequenceNumber()}; @@ -140,9 +137,8 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { SpecialFields::RowId().Name(), SpecialFields::SequenceNumber().Name(), }; - ASSERT_OK_AND_ASSIGN(auto result, - AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols)); + ASSERT_OK_AND_ASSIGN(auto result, AbstractSplitRead::BuildDataFieldsForFieldMapping( + table_schema, write_cols)); std::vector expected = {fields[3], fields[0], fields[1], SpecialFields::RowId(), SpecialFields::SequenceNumber()}; @@ -151,9 +147,9 @@ TEST(AbstractSplitReadTest, TestProjectFieldsForRowTrackingAndDataEvolution) { { // test with empty write_cols std::vector write_cols = {}; - ASSERT_NOK_WITH_MSG(AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution( - table_schema, write_cols), - "write cols cannot be empty"); + ASSERT_NOK_WITH_MSG( + AbstractSplitRead::BuildDataFieldsForFieldMapping(table_schema, write_cols), + "write cols cannot be empty"); } } } 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 135c49dc..2f4073af 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -27,6 +27,7 @@ #include "paimon/common/data/binary_row.h" #include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/core/append/append_only_writer.h" #include "paimon/core/append/bucketed_append_compact_manager.h" @@ -40,10 +41,12 @@ #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/operation/append_only_file_store_scan.h" +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/operation/raw_file_split_read.h" #include "paimon/core/operation/restore_files.h" +#include "paimon/core/realtime/realtime_append_only_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" #include "paimon/core/utils/file_store_path_factory.h" @@ -51,6 +54,7 @@ #include "paimon/executor.h" #include "paimon/logging.h" #include "paimon/read_context.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/result.h" namespace arrow { class Schema; @@ -72,11 +76,15 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, + const std::shared_ptr& realtime_context, + const RealtimeSnapshotProperties::OffsetMap& realtime_committed_offsets, 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), + realtime_context_(realtime_context), + realtime_committed_offsets_(realtime_committed_offsets), logger_(Logger::GetLogger("AppendOnlyFileStoreWrite")) { write_cols_ = write_schema->field_names(); // optimize write_cols to null in following cases: @@ -86,6 +94,13 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( if (schema->Equals(write_schema)) { write_cols_ = std::nullopt; } + if (realtime_context_) { + writer_memory_manager_ = std::make_unique(); + arrow::FieldVector fields = write_schema->fields(); + fields.insert(fields.begin(), + DataField::ConvertDataFieldToArrowField(SpecialFields::Offset())); + realtime_write_schema_ = arrow::schema(fields); + } } AppendOnlyFileStoreWrite::~AppendOnlyFileStoreWrite() = default; @@ -182,7 +197,8 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( file_store_path_factory_->CreateDataFilePathFactory(partition, bucket)); std::shared_ptr compact_manager; - if (options_.WriteOnly() || options_.DataEvolutionEnabled() || options_.GetBucket() == -1) { + if (realtime_context_ || options_.WriteOnly() || options_.DataEvolutionEnabled() || + options_.GetBucket() == -1) { compact_manager = std::make_shared(); } else { auto dv_factory = @@ -212,10 +228,32 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( compaction_metrics_->CreateReporter(partition, bucket), cancellation_controller); } - auto writer = std::make_unique( - options_, table_schema_->Id(), write_schema_, write_cols_, restore_max_seq_number, + const std::shared_ptr& file_write_schema = + realtime_context_ ? realtime_write_schema_ : write_schema_; + auto writer = std::make_shared( + options_, table_schema_->Id(), file_write_schema, write_cols_, restore_max_seq_number, data_file_path_factory, compact_manager, pool_); - return std::shared_ptr(std::move(writer)); + if (!realtime_context_) { + return std::shared_ptr(std::move(writer)); + } + + PAIMON_ASSIGN_OR_RAISE(std::string partition_string, + file_store_path_factory_->GetPartitionString(partition)); + partition_string = PartitionBucket::NormalizePartition(std::move(partition_string)); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*write_schema_, c_write_schema.get())); + int64_t next_offset = 0; + PartitionBucket partition_bucket(partition_string, bucket); + auto offset_iter = realtime_committed_offsets_.find(partition_bucket); + if (offset_iter != realtime_committed_offsets_.end()) { + if (offset_iter->second == std::numeric_limits::max()) { + return Status::Invalid("real-time offset has reached INT64_MAX"); + } + next_offset = offset_iter->second + 1; + } + return RealtimeAppendOnlyWriter::Create( + partition_string, bucket, std::move(c_write_schema), realtime_context_, writer, + write_schema_, realtime_write_schema_, options_.ToMap(), next_offset, 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 905a4180..442af906 100644 --- a/src/paimon/core/operation/append_only_file_store_write.h +++ b/src/paimon/core/operation/append_only_file_store_write.h @@ -32,6 +32,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/deletion_vector.h" #include "paimon/core/operation/abstract_file_store_write.h" +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" #include "paimon/core/table/bucket_mode.h" #include "paimon/file_store_write.h" #include "paimon/logging.h" @@ -81,6 +82,8 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr& dv_maintainer_factory, const std::shared_ptr& io_manager, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, + const std::shared_ptr& realtime_context, + const RealtimeSnapshotProperties::OffsetMap& realtime_committed_offsets, const std::shared_ptr& executor, const std::shared_ptr& pool); ~AppendOnlyFileStoreWrite() override; @@ -110,6 +113,10 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { Result> CreateFileStoreScan( const std::shared_ptr& filter) const override; + bool IsRealtimeWrite() const override { + return realtime_context_ != nullptr; + } + Result GetDataFileWriterFactory( const std::shared_ptr& data_file_path_factory, const std::shared_ptr& schema, @@ -121,6 +128,9 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { const std::vector>& files) const; std::optional> write_cols_; + std::shared_ptr realtime_write_schema_; + std::shared_ptr realtime_context_; + RealtimeSnapshotProperties::OffsetMap realtime_committed_offsets_; std::unique_ptr logger_; }; 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 df1e1bd2..07b2d9cb 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 @@ -26,6 +26,7 @@ #include "arrow/array/array_base.h" #include "arrow/array/builder_binary.h" +#include "arrow/array/concatenate.h" #include "arrow/c/abi.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" @@ -54,8 +55,10 @@ #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" +#include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" #include "paimon/write_context.h" @@ -169,6 +172,31 @@ class AppendOnlyFileStoreWriteTest : public testing::Test { return arrow::ImportSchema(c_file_schema.get()).ValueOrDie(); } + std::shared_ptr ReadDataFileArray( + const std::string& table_path, const std::shared_ptr& file, + const std::map& options) const { + std::string file_path = + PathUtil::JoinPath(PathUtil::JoinPath(table_path, "bucket-0"), file->file_name); + auto fs = std::make_shared(); + EXPECT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs->Open(file_path)); + EXPECT_OK_AND_ASSIGN(auto format_str, file->FileFormat()); + EXPECT_OK_AND_ASSIGN(auto file_format, FileFormatFactory::Get(format_str, options)); + EXPECT_OK_AND_ASSIGN(auto reader_builder, file_format->CreateReaderBuilder(10)); + EXPECT_OK_AND_ASSIGN(auto reader, reader_builder->Build(input_stream)); + EXPECT_OK_AND_ASSIGN(auto c_file_schema, reader->GetFileSchema()); + EXPECT_OK(reader->SetReadSchema(c_file_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + EXPECT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(reader.get())); + EXPECT_NE(nullptr, result); + if (!result) { + return nullptr; + } + std::shared_ptr array = + arrow::Concatenate(result->chunks(), arrow::default_memory_pool()).ValueOrDie(); + return std::static_pointer_cast(array); + } + MapSharedShreddingFieldMeta ShreddingMeta(const std::shared_ptr& file_schema, int32_t field_index) const { auto metadata = file_schema->field(field_index)->metadata(); @@ -199,6 +227,8 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestWriteWithInvalidBatch) { ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); + ASSERT_NOK_WITH_MSG(file_store_write->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PrepareCommitWithProgress is only supported by a real-time writer"); ASSERT_NOK_WITH_MSG(file_store_write->Write(nullptr), "batch is null pointer"); } { @@ -236,6 +266,93 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestWriteWithInvalidBatch) { } } +TEST_F(AppendOnlyFileStoreWriteTest, TestRealtimeWriteFlushesOffsetColumn) { + std::map options = { + {"file.format", "parquet"}, + {"write-only", "true"}, + {"bucket", "1"}, + {"bucket-key", "id"}, + }; + auto logical_schema = + arrow::schema({arrow::field("id", arrow::int32()), arrow::field("name", arrow::utf8())}); + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), logical_schema, options); + std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + { + ASSERT_OK_AND_ASSIGN(std::shared_ptr batch_realtime_context, + RealtimeContext::Create()); + WriteContextBuilder batch_builder(table_path, commit_user_); + batch_builder.SetOptions(options).WithRealtimeContext(batch_realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_context, batch_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_writer, + FileStoreWrite::Create(std::move(batch_context))); + ASSERT_NOK_WITH_MSG(batch_writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PrepareCommitWithProgress requires streaming mode"); + ASSERT_OK(batch_writer->Close()); + } + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, commit_user_); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto file_store_write, FileStoreWrite::Create(std::move(write_context))); + + ASSERT_OK(file_store_write->Write(MakeBatch(logical_schema, R"([ + [1, "a"], + [2, "b"] + ])"))); + ASSERT_NOK_WITH_MSG( + file_store_write->PrepareCommit(/*wait_compaction=*/false, /*commit_identifier=*/0), + "real-time writer must use PrepareCommitWithProgress"); + ASSERT_OK_AND_ASSIGN(auto first_prepared, + file_store_write->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, first_prepared.size()); + ASSERT_EQ("", first_prepared[0].partition); + ASSERT_EQ(0, first_prepared[0].bucket); + ASSERT_EQ(Range(0, 1), first_prepared[0].offset_range); + std::shared_ptr first_file = OnlyNewFile({first_prepared[0].commit_message}); + ASSERT_FALSE(first_file->write_cols.has_value()); + std::shared_ptr first_schema = + ReadDataFileSchema(table_path, first_file, options); + ASSERT_EQ(3, first_schema->num_fields()); + ASSERT_EQ("_OFFSET", first_schema->field(0)->name()); + ASSERT_EQ(arrow::Type::INT64, first_schema->field(0)->type()->id()); + std::shared_ptr first_array = + ReadDataFileArray(table_path, first_file, options); + ASSERT_NE(nullptr, first_array); + std::shared_ptr expected_first_array = + arrow::ipc::internal::json::ArrayFromJSON(first_array->type(), R"([ + [0, 1, "a"], + [1, 2, "b"] + ])") + .ValueOrDie(); + ASSERT_TRUE(first_array->Equals(*expected_first_array)) << first_array->ToString(); + + ASSERT_OK(file_store_write->Write(MakeBatch(logical_schema, R"([ + [3, "c"] + ])"))); + ASSERT_OK_AND_ASSIGN(auto second_prepared, + file_store_write->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_prepared.size()); + ASSERT_EQ("", second_prepared[0].partition); + ASSERT_EQ(0, second_prepared[0].bucket); + ASSERT_EQ(Range(2, 2), second_prepared[0].offset_range); + std::shared_ptr second_file = OnlyNewFile({second_prepared[0].commit_message}); + std::shared_ptr second_array = + ReadDataFileArray(table_path, second_file, options); + ASSERT_NE(nullptr, second_array); + std::shared_ptr expected_second_array = + arrow::ipc::internal::json::ArrayFromJSON(second_array->type(), R"([ + [2, 3, "c"] + ])") + .ValueOrDie(); + ASSERT_TRUE(second_array->Equals(*expected_second_array)) << second_array->ToString(); + ASSERT_OK(file_store_write->Close()); +} + TEST_F(AppendOnlyFileStoreWriteTest, TestGetMaxSequenceNumberFromMultiPartition) { WriteContextBuilder builder( paimon::test::GetDataDir() + diff --git a/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp b/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp new file mode 100644 index 00000000..b4b6f30f --- /dev/null +++ b/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp @@ -0,0 +1,287 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "fmt/format.h" +#include "paimon/common/utils/rapidjson_util.h" +#include "paimon/common/utils/uuid.h" +#include "paimon/core/utils/branch_manager.h" +#include "paimon/fs/file_system.h" +#include "paimon/macros.h" + +namespace paimon { +namespace { + +class OffsetEntryJson { + public: + OffsetEntryJson() = default; + + OffsetEntryJson(std::string partition, int32_t bucket, int64_t offset) + : partition_(std::move(partition)), bucket_(bucket), offset_(offset) {} + + rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const { + rapidjson::Value value(rapidjson::kObjectType); + value.AddMember("partition", RapidJsonUtil::SerializeValue(partition_, allocator), + *allocator); + value.AddMember("bucket", RapidJsonUtil::SerializeValue(bucket_, allocator), *allocator); + value.AddMember("offset", RapidJsonUtil::SerializeValue(offset_, allocator), *allocator); + return value; + } + + void FromJson(const rapidjson::Value& value) { + partition_ = RapidJsonUtil::DeserializeKeyValue(value, "partition"); + bucket_ = RapidJsonUtil::DeserializeKeyValue(value, "bucket"); + offset_ = RapidJsonUtil::DeserializeKeyValue(value, "offset"); + } + + const std::string& Partition() const { + return partition_; + } + + int32_t Bucket() const { + return bucket_; + } + + int64_t Offset() const { + return offset_; + } + + private: + std::string partition_; + int32_t bucket_ = -1; + int64_t offset_ = -1; +}; + +class OffsetsJson { + public: + OffsetsJson() = default; + + explicit OffsetsJson(const RealtimeSnapshotProperties::OffsetMap& offsets) + : offsets_(offsets) {} + + rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const { + rapidjson::Value value(rapidjson::kObjectType); + value.AddMember( + "version", + RapidJsonUtil::SerializeValue(RealtimeSnapshotProperties::kOffsetsVersion, allocator), + *allocator); + std::vector entries; + entries.reserve(offsets_.size()); + for (const auto& [partition_bucket, offset] : offsets_) { + if (partition_bucket.bucket < 0) { + throw std::invalid_argument( + fmt::format("invalid bucket {} in offsets", partition_bucket.bucket)); + } + if (offset < 0) { + throw std::invalid_argument( + fmt::format("invalid offset {} for partition '{}' " + "bucket {}", + offset, partition_bucket.partition, partition_bucket.bucket)); + } + entries.emplace_back(partition_bucket.partition, partition_bucket.bucket, offset); + } + value.AddMember("offsets", RapidJsonUtil::SerializeValue(entries, allocator), *allocator); + return value; + } + + void FromJson(const rapidjson::Value& value) { + int32_t version = RapidJsonUtil::DeserializeKeyValue(value, "version"); + if (version != RealtimeSnapshotProperties::kOffsetsVersion) { + throw std::invalid_argument(fmt::format("unsupported offsets version {}", version)); + } + std::vector entries = + RapidJsonUtil::DeserializeKeyValue>(value, "offsets"); + offsets_.clear(); + for (const OffsetEntryJson& entry : entries) { + if (entry.Bucket() < 0) { + throw std::invalid_argument( + fmt::format("invalid bucket {} in offsets", entry.Bucket())); + } + if (entry.Offset() < 0) { + throw std::invalid_argument( + fmt::format("invalid offset {} in offsets", entry.Offset())); + } + PartitionBucket partition_bucket(entry.Partition(), entry.Bucket()); + if (!offsets_.emplace(std::move(partition_bucket), entry.Offset()).second) { + throw std::invalid_argument( + fmt::format("duplicate partition '{}' bucket {} in offsets", entry.Partition(), + entry.Bucket())); + } + } + } + + const RealtimeSnapshotProperties::OffsetMap& Offsets() const { + return offsets_; + } + + private: + RealtimeSnapshotProperties::OffsetMap offsets_; +}; + +} // namespace + +Result +RealtimeSnapshotProperties::ValidateProgress(const std::vector& commits, + const OffsetMap& committed_offsets) { + ValidatedCommitProgress result; + result.ordered_commits = commits; + std::stable_sort(result.ordered_commits.begin(), result.ordered_commits.end(), + [](const RealtimeCommitProgress& lhs, const RealtimeCommitProgress& rhs) { + if (lhs.partition != rhs.partition) { + return lhs.partition < rhs.partition; + } + if (lhs.bucket != rhs.bucket) { + return lhs.bucket < rhs.bucket; + } + return lhs.offset_range.from < rhs.offset_range.from; + }); + + OffsetMap last_offsets; + for (const RealtimeCommitProgress& commit : result.ordered_commits) { + if (commit.bucket < 0) { + return Status::Invalid( + fmt::format("real-time commit bucket {} is invalid", commit.bucket)); + } + if (commit.offset_range.from > commit.offset_range.to) { + return Status::Invalid("real-time commit offset range is invalid"); + } + + PartitionBucket partition_bucket(commit.partition, commit.bucket); + auto last_iter = last_offsets.find(partition_bucket); + int64_t previous_offset = -1; + if (last_iter != last_offsets.end()) { + previous_offset = last_iter->second; + } else { + auto committed_iter = committed_offsets.find(partition_bucket); + if (committed_iter != committed_offsets.end()) { + previous_offset = committed_iter->second; + } + } + if (previous_offset == std::numeric_limits::max() || + commit.offset_range.from != previous_offset + 1) { + return Status::Invalid(fmt::format( + "real-time commit offsets for partition '{}' bucket {} are not contiguous", + commit.partition, commit.bucket)); + } + + last_offsets[partition_bucket] = commit.offset_range.to; + result.delta_offsets[partition_bucket] = commit.offset_range.to; + } + return result; +} + +std::string RealtimeSnapshotProperties::OffsetsDirectory(const std::string& table_root, + const std::string& branch) { + return PathUtil::JoinPath(BranchManager::BranchPath(table_root, branch), "metadata"); +} + +Result RealtimeSnapshotProperties::ReadOffsets( + const std::optional& snapshot, const std::shared_ptr& file_system) { + if (!snapshot || !snapshot->Properties()) { + return OffsetMap{}; + } + const std::map& properties = snapshot->Properties().value(); + auto iter = properties.find(kOffsetsKey); + if (iter == properties.end()) { + return OffsetMap{}; + } + if (file_system == nullptr) { + return Status::Invalid("file system is null when reading real-time offsets"); + } + std::string content; + PAIMON_RETURN_NOT_OK(file_system->ReadFile(iter->second, &content)); + return ParseOffsets(content); +} + +Result RealtimeSnapshotProperties::SerializeOffsets(const OffsetMap& offsets) { + std::string result; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(OffsetsJson(offsets), &result)); + return result; +} + +Result> RealtimeSnapshotProperties::MergeOffsets( + const std::map& properties, + const std::optional& latest_snapshot, const std::shared_ptr& file_system, + const std::string& offsets_directory) { + std::map merged_properties = properties; + auto delta_iter = properties.find(kOffsetsDeltaKey); + merged_properties.erase(kOffsetsDeltaKey); + if (delta_iter == properties.end()) { + if (latest_snapshot && latest_snapshot->Properties()) { + const std::map& latest_properties = + latest_snapshot->Properties().value(); + auto offsets_iter = latest_properties.find(kOffsetsKey); + if (offsets_iter != latest_properties.end()) { + merged_properties[kOffsetsKey] = offsets_iter->second; + } + } + return merged_properties; + } + + PAIMON_ASSIGN_OR_RAISE(OffsetMap merged_offsets, ReadOffsets(latest_snapshot, file_system)); + PAIMON_ASSIGN_OR_RAISE(OffsetMap delta_offsets, ParseOffsets(delta_iter->second)); + for (const auto& [key, offset] : delta_offsets) { + auto merged_iter = merged_offsets.find(key); + if (merged_iter == merged_offsets.end()) { + merged_offsets.emplace(key, offset); + } else { + merged_iter->second = std::max(merged_iter->second, offset); + } + } + if (!merged_offsets.empty()) { + PAIMON_ASSIGN_OR_RAISE(merged_properties[kOffsetsKey], + WriteOffsets(merged_offsets, file_system, offsets_directory)); + } + return merged_properties; +} + +Result RealtimeSnapshotProperties::WriteOffsets( + const OffsetMap& offsets, const std::shared_ptr& file_system, + const std::string& offsets_directory) { + if (file_system == nullptr) { + return Status::Invalid("file system is null when writing real-time offsets"); + } + if (offsets_directory.empty()) { + return Status::Invalid("real-time offsets directory is empty"); + } + std::string uuid; + if (!UUID::Generate(&uuid)) { + return Status::Invalid("fail to generate uuid for real-time offsets file"); + } + PAIMON_RETURN_NOT_OK(file_system->Mkdirs(offsets_directory)); + std::string path = PathUtil::JoinPath(offsets_directory, uuid + ".offsets"); + PAIMON_ASSIGN_OR_RAISE(std::string content, SerializeOffsets(offsets)); + PAIMON_RETURN_NOT_OK(file_system->WriteFile(path, content, /*overwrite=*/false)); + return path; +} + +Result RealtimeSnapshotProperties::ParseOffsets( + const std::string& value) { + OffsetsJson offsets_json; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::FromJsonString(value, &offsets_json)); + return offsets_json.Offsets(); +} + +} // namespace paimon diff --git a/src/paimon/core/operation/commit/realtime_snapshot_properties.h b/src/paimon/core/operation/commit/realtime_snapshot_properties.h new file mode 100644 index 00000000..64b4f15d --- /dev/null +++ b/src/paimon/core/operation/commit/realtime_snapshot_properties.h @@ -0,0 +1,75 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/core/realtime/partition_bucket.h" +#include "paimon/core/snapshot.h" +#include "paimon/realtime/realtime_commit_progress.h" +#include "paimon/result.h" + +namespace paimon { + +class FileSystem; + +class RealtimeSnapshotProperties { + public: + using OffsetMap = std::map; + + /// Ordered commit progress and the latest offset contributed by each partition-bucket. + struct ValidatedCommitProgress { + std::vector ordered_commits; + OffsetMap delta_offsets; + }; + + RealtimeSnapshotProperties() = delete; + + static constexpr const char* kOffsetsKey = "realtime.offsets"; + static constexpr const char* kOffsetsDeltaKey = "realtime.offsets.delta"; + static constexpr int32_t kOffsetsVersion = 1; + + /// Orders progress entries and verifies they form a contiguous prefix after committed offsets. + static Result ValidateProgress( + const std::vector& commits, const OffsetMap& committed_offsets); + + static std::string OffsetsDirectory(const std::string& table_root, const std::string& branch); + + static Result ReadOffsets(const std::optional& snapshot, + const std::shared_ptr& file_system); + + static Result SerializeOffsets(const OffsetMap& offsets); + + static Result> MergeOffsets( + const std::map& properties, + const std::optional& latest_snapshot, + const std::shared_ptr& file_system, const std::string& offsets_directory); + + private: + static Result WriteOffsets(const OffsetMap& offsets, + const std::shared_ptr& file_system, + const std::string& offsets_directory); + + static Result ParseOffsets(const std::string& value); +}; + +} // namespace paimon diff --git a/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp b/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp new file mode 100644 index 00000000..512e64ca --- /dev/null +++ b/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp @@ -0,0 +1,363 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "paimon/fs/file_system.h" +#include "paimon/macros.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +using Properties = std::map; + +const std::string kTargetJson = R"({ + "version": 1, + "offsets": [ + { + "partition": "", + "bucket": 0, + "offset": 7 + }, + { + "partition": "dt=2", + "bucket": 2, + "offset": 9 + }, + { + "partition": "dt=2", + "bucket": 10, + "offset": 12 + } + ] +})"; + +class RealtimeSnapshotPropertiesTest : public testing::Test { + protected: + void SetUp() override { + directory_ = UniqueTestDirectory::Create(); + ASSERT_NE(nullptr, directory_); + file_system_ = directory_->GetFileSystem(); + ASSERT_NE(nullptr, file_system_); + } + + Snapshot MakeSnapshot( + const std::optional>& properties) const { + return Snapshot( + /*id=*/1, + /*schema_id=*/1, + /*base_manifest_list=*/"base-manifest-list", + /*base_manifest_list_size=*/std::nullopt, + /*delta_manifest_list=*/"delta-manifest-list", + /*delta_manifest_list_size=*/std::nullopt, + /*changelog_manifest_list=*/std::nullopt, + /*changelog_manifest_list_size=*/std::nullopt, + /*index_manifest=*/std::nullopt, + /*commit_user=*/"test-user", + /*commit_identifier=*/1, Snapshot::CommitKind::Append(), + /*time_millis=*/0, + /*total_record_count=*/0, + /*delta_record_count=*/0, + /*changelog_record_count=*/std::nullopt, + /*watermark=*/std::nullopt, + /*statistics=*/std::nullopt, properties, + /*next_row_id=*/std::nullopt); + } + + Result WriteJson(const std::string& json) { + std::string path = + directory_->Str() + "/input-" + std::to_string(next_file_id_++) + ".offsets"; + PAIMON_RETURN_NOT_OK(file_system_->WriteFile(path, json, /*overwrite=*/false)); + return path; + } + + Result ReadJson(const std::string& json) { + PAIMON_ASSIGN_OR_RAISE(std::string path, WriteJson(json)); + std::map properties = { + {RealtimeSnapshotProperties::kOffsetsKey, path}}; + return RealtimeSnapshotProperties::ReadOffsets( + std::optional(MakeSnapshot(properties)), file_system_); + } + + std::unique_ptr directory_; + std::shared_ptr file_system_; + int32_t next_file_id_ = 0; +}; + +RealtimeSnapshotProperties::OffsetMap TargetOffsets() { + return {{PartitionBucket("", /*bucket=*/0), 7}, + {PartitionBucket("dt=2", /*bucket=*/2), 9}, + {PartitionBucket("dt=2", /*bucket=*/10), 12}}; +} + +} // namespace + +TEST_F(RealtimeSnapshotPropertiesTest, SerializeAndDeserializeTargetJson) { + RealtimeSnapshotProperties::OffsetMap expected = TargetOffsets(); + ASSERT_OK_AND_ASSIGN(std::string actual_json, + RealtimeSnapshotProperties::SerializeOffsets(expected)); + ASSERT_EQ(kTargetJson, actual_json); + + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap actual, ReadJson(kTargetJson)); + ASSERT_EQ(expected, actual); + ASSERT_OK_AND_ASSIGN(std::string round_trip_json, + RealtimeSnapshotProperties::SerializeOffsets(actual)); + ASSERT_EQ(kTargetJson, round_trip_json); +} + +TEST_F(RealtimeSnapshotPropertiesTest, SerializeEmptyOffsets) { + ASSERT_OK_AND_ASSIGN(std::string actual_json, + RealtimeSnapshotProperties::SerializeOffsets(/*offsets=*/{})); + ASSERT_EQ(R"({ + "version": 1, + "offsets": [] +})", + actual_json); +} + +TEST_F(RealtimeSnapshotPropertiesTest, SerializeRejectsInvalidOffsets) { + RealtimeSnapshotProperties::OffsetMap negative_bucket = {{PartitionBucket("dt=2", -1), 0}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::SerializeOffsets(negative_bucket), + "invalid bucket -1"); + + RealtimeSnapshotProperties::OffsetMap negative_offset = { + {PartitionBucket("dt=2", /*bucket=*/0), -1}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::SerializeOffsets(negative_offset), + "invalid offset -1"); +} + +TEST_F(RealtimeSnapshotPropertiesTest, NormalizePartitionBucketAndOffsetsDirectory) { + ASSERT_EQ("dt=2", PartitionBucket::NormalizePartition("dt=2/")); + PartitionBucket expected_partition_bucket("dt=2", 3); + ASSERT_EQ(expected_partition_bucket, PartitionBucket("dt=2/", /*bucket=*/3)); + ASSERT_EQ("/table/metadata", RealtimeSnapshotProperties::OffsetsDirectory("/table", "main")); + ASSERT_EQ("/table/branch/branch-dev/metadata", + RealtimeSnapshotProperties::OffsetsDirectory("/table", "dev")); +} + +TEST_F(RealtimeSnapshotPropertiesTest, ReadOffsetsWithoutProgress) { + ASSERT_OK_AND_ASSIGN( + RealtimeSnapshotProperties::OffsetMap no_snapshot, + RealtimeSnapshotProperties::ReadOffsets(/*snapshot=*/std::nullopt, file_system_)); + ASSERT_TRUE(no_snapshot.empty()); + + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap no_properties, + RealtimeSnapshotProperties::ReadOffsets( + std::optional(MakeSnapshot(std::nullopt)), file_system_)); + ASSERT_TRUE(no_properties.empty()); + + std::map unrelated_properties = {{"other", "value"}}; + ASSERT_OK_AND_ASSIGN( + RealtimeSnapshotProperties::OffsetMap no_offsets_property, + RealtimeSnapshotProperties::ReadOffsets( + std::optional(MakeSnapshot(unrelated_properties)), file_system_)); + ASSERT_TRUE(no_offsets_property.empty()); +} + +TEST_F(RealtimeSnapshotPropertiesTest, ReadOffsetsRequiresFileSystem) { + std::map properties = { + {RealtimeSnapshotProperties::kOffsetsKey, "metadata/test.offsets"}}; + ASSERT_NOK_WITH_MSG( + RealtimeSnapshotProperties::ReadOffsets(std::optional(MakeSnapshot(properties)), + /*file_system=*/nullptr), + "file system is null"); +} + +TEST_F(RealtimeSnapshotPropertiesTest, ReadOffsetsPropagatesFileError) { + std::map properties = { + {RealtimeSnapshotProperties::kOffsetsKey, directory_->Str() + "/missing.offsets"}}; + ASSERT_NOK(RealtimeSnapshotProperties::ReadOffsets( + std::optional(MakeSnapshot(properties)), file_system_)); +} + +TEST_F(RealtimeSnapshotPropertiesTest, DeserializeRejectsInvalidJson) { + const std::vector> invalid_json_cases = { + {"{", "deserialize failed"}, + {R"([])", "value must be an object"}, + {R"({"offsets":[]})", "key must exist"}, + {R"({"version":2,"offsets":[]})", "unsupported offsets version 2"}, + {R"({"version":1,"offsets":{}})", "value must be an array"}, + {R"({"version":1,"offsets":[{"bucket":0,"offset":1}]})", "key must exist"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","offset":1}]})", "key must exist"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0}]})", "key must exist"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":"0","offset":1}]})", + "value must be int"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":"1"}]})", + "value must be int64"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":-1,"offset":1}]})", + "invalid bucket -1"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":-1}]})", + "invalid offset -1"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":1},{"partition":"dt=2/","bucket":0,"offset":2}]})", + "duplicate partition 'dt=2/' bucket 0"}}; + + for (const auto& [json, expected_error] : invalid_json_cases) { + SCOPED_TRACE(json); + ASSERT_NOK_WITH_MSG(ReadJson(json), expected_error); + } +} + +TEST_F(RealtimeSnapshotPropertiesTest, ValidateAndOrderProgress) { + PartitionBucket bucket0("dt=2", /*bucket=*/0); + PartitionBucket bucket1("dt=2", /*bucket=*/1); + RealtimeSnapshotProperties::OffsetMap committed_offsets = {{bucket1, 4}}; + std::vector commits = { + {/*commit_message=*/nullptr, "dt=2", 0, Range(2, 3)}, + {/*commit_message=*/nullptr, "dt=2", 1, Range(5, 6)}, + {/*commit_message=*/nullptr, "dt=2", 0, Range(0, 1)}}; + + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::ValidatedCommitProgress validated_progress, + RealtimeSnapshotProperties::ValidateProgress(commits, committed_offsets)); + ASSERT_EQ(3, validated_progress.ordered_commits.size()); + ASSERT_EQ(Range(0, 1), validated_progress.ordered_commits[0].offset_range); + ASSERT_EQ(Range(2, 3), validated_progress.ordered_commits[1].offset_range); + ASSERT_EQ(Range(5, 6), validated_progress.ordered_commits[2].offset_range); + ASSERT_EQ(3, validated_progress.delta_offsets.at(bucket0)); + ASSERT_EQ(6, validated_progress.delta_offsets.at(bucket1)); +} + +TEST_F(RealtimeSnapshotPropertiesTest, ValidateProgressRejectsInvalidProgress) { + PartitionBucket bucket0("dt=2", /*bucket=*/0); + RealtimeSnapshotProperties::OffsetMap committed_offsets = {{bucket0, 1}}; + + std::vector invalid_bucket = { + {/*commit_message=*/nullptr, "dt=2", -1, Range(0, 0)}}; + ASSERT_NOK_WITH_MSG( + RealtimeSnapshotProperties::ValidateProgress(invalid_bucket, /*committed_offsets=*/{}), + "bucket -1 is invalid"); + + std::vector gap = { + {/*commit_message=*/nullptr, "dt=2", 0, Range(3, 4)}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(gap, committed_offsets), + "are not contiguous"); + + std::vector overlap = { + {/*commit_message=*/nullptr, "dt=2", 0, Range(1, 2)}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(overlap, committed_offsets), + "are not contiguous"); + + RealtimeSnapshotProperties::OffsetMap exhausted_offsets = { + {bucket0, std::numeric_limits::max()}}; + std::vector after_max = { + {/*commit_message=*/nullptr, "dt=2", 0, + Range(std::numeric_limits::max(), std::numeric_limits::max())}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(after_max, exhausted_offsets), + "are not contiguous"); +} + +TEST_F(RealtimeSnapshotPropertiesTest, ValidateEmptyProgress) { + ASSERT_OK_AND_ASSIGN( + RealtimeSnapshotProperties::ValidatedCommitProgress validated_progress, + RealtimeSnapshotProperties::ValidateProgress(/*commits=*/{}, /*committed_offsets=*/{})); + ASSERT_TRUE(validated_progress.ordered_commits.empty()); + ASSERT_TRUE(validated_progress.delta_offsets.empty()); +} + +TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsWithoutDelta) { + std::string latest_offsets_path = directory_->Str() + "/latest.offsets"; + std::map latest_properties = { + {RealtimeSnapshotProperties::kOffsetsKey, latest_offsets_path}}; + std::map properties = {{"custom", "value"}}; + + ASSERT_OK_AND_ASSIGN(Properties merged, + RealtimeSnapshotProperties::MergeOffsets( + properties, std::optional(MakeSnapshot(latest_properties)), + file_system_, directory_->Str() + "/metadata")); + ASSERT_EQ("value", merged.at("custom")); + ASSERT_EQ(latest_offsets_path, merged.at(RealtimeSnapshotProperties::kOffsetsKey)); + + ASSERT_OK_AND_ASSIGN( + Properties merged_without_snapshot, + RealtimeSnapshotProperties::MergeOffsets(properties, /*latest_snapshot=*/std::nullopt, + file_system_, directory_->Str() + "/metadata")); + ASSERT_EQ(properties, merged_without_snapshot); +} + +TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsWritesMergedProgress) { + ASSERT_OK_AND_ASSIGN(std::string latest_offsets_path, WriteJson(kTargetJson)); + std::map latest_properties = { + {RealtimeSnapshotProperties::kOffsetsKey, latest_offsets_path}}; + + RealtimeSnapshotProperties::OffsetMap delta_offsets = { + {PartitionBucket("", /*bucket=*/0), 5}, + {PartitionBucket("dt=2", /*bucket=*/2), 11}, + {PartitionBucket("dt=3", /*bucket=*/0), 4}}; + ASSERT_OK_AND_ASSIGN(std::string delta_json, + RealtimeSnapshotProperties::SerializeOffsets(delta_offsets)); + std::map properties = { + {"custom", "value"}, {RealtimeSnapshotProperties::kOffsetsDeltaKey, delta_json}}; + + ASSERT_OK_AND_ASSIGN(Properties merged, + RealtimeSnapshotProperties::MergeOffsets( + properties, std::optional(MakeSnapshot(latest_properties)), + file_system_, directory_->Str() + "/metadata")); + ASSERT_EQ("value", merged.at("custom")); + ASSERT_EQ(0, merged.count(RealtimeSnapshotProperties::kOffsetsDeltaKey)); + ASSERT_NE(latest_offsets_path, merged.at(RealtimeSnapshotProperties::kOffsetsKey)); + + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap actual, + RealtimeSnapshotProperties::ReadOffsets( + std::optional(MakeSnapshot(merged)), file_system_)); + RealtimeSnapshotProperties::OffsetMap expected = TargetOffsets(); + expected[PartitionBucket("dt=2", /*bucket=*/2)] = 11; + expected[PartitionBucket("dt=3", /*bucket=*/0)] = 4; + ASSERT_EQ(expected, actual); +} + +TEST_F(RealtimeSnapshotPropertiesTest, MergeEmptyDelta) { + ASSERT_OK_AND_ASSIGN(std::string empty_delta, + RealtimeSnapshotProperties::SerializeOffsets(/*offsets=*/{})); + std::map properties = { + {"custom", "value"}, {RealtimeSnapshotProperties::kOffsetsDeltaKey, empty_delta}}; + + ASSERT_OK_AND_ASSIGN(Properties merged, + RealtimeSnapshotProperties::MergeOffsets( + properties, /*latest_snapshot=*/std::nullopt, /*file_system=*/nullptr, + /*offsets_directory=*/"")); + std::map expected = {{"custom", "value"}}; + ASSERT_EQ(expected, merged); +} + +TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsRequiresFileSystem) { + RealtimeSnapshotProperties::OffsetMap delta_offsets = { + {PartitionBucket("dt=2", /*bucket=*/0), 1}}; + ASSERT_OK_AND_ASSIGN(std::string delta_json, + RealtimeSnapshotProperties::SerializeOffsets(delta_offsets)); + std::map properties = { + {RealtimeSnapshotProperties::kOffsetsDeltaKey, delta_json}}; + ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::MergeOffsets( + properties, /*latest_snapshot=*/std::nullopt, /*file_system=*/nullptr, + directory_->Str() + "/metadata"), + "file system is null"); + + ASSERT_NOK_WITH_MSG( + RealtimeSnapshotProperties::MergeOffsets(properties, /*latest_snapshot=*/std::nullopt, + file_system_, /*offsets_directory=*/""), + "offsets directory is empty"); +} + +} // namespace paimon::test diff --git a/src/paimon/core/operation/data_evolution_split_read.cpp b/src/paimon/core/operation/data_evolution_split_read.cpp index 033d3e07..1f4f1937 100644 --- a/src/paimon/core/operation/data_evolution_split_read.cpp +++ b/src/paimon/core/operation/data_evolution_split_read.cpp @@ -458,9 +458,8 @@ Result> DataEvolutionSplitRead::CreateU if (schema_id != data_schema->Id()) { PAIMON_ASSIGN_OR_RAISE(data_schema, schema_manager_->ReadSchema(schema_id)); } - PAIMON_ASSIGN_OR_RAISE( - std::vector data_fields, - ProjectFieldsForRowTrackingAndDataEvolution(data_schema, first_file->write_cols)); + PAIMON_ASSIGN_OR_RAISE(std::vector data_fields, + BuildDataFieldsForFieldMapping(data_schema, first_file->write_cols)); std::vector data_field_ids = DataField::GetAllFieldIds(data_fields); std::vector read_fields_in_file; for (int32_t read_field_idx = 0; diff --git a/src/paimon/core/operation/file_store_commit_impl.cpp b/src/paimon/core/operation/file_store_commit_impl.cpp index 6342f9b3..3c58979e 100644 --- a/src/paimon/core/operation/file_store_commit_impl.cpp +++ b/src/paimon/core/operation/file_store_commit_impl.cpp @@ -66,6 +66,7 @@ #include "paimon/core/operation/commit/commit_changes_provider.h" #include "paimon/core/operation/commit/compacted_changelog_path_resolver.h" #include "paimon/core/operation/commit/conflict_detection.h" +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" #include "paimon/core/operation/commit/row_tracking_commit_utils.h" #include "paimon/core/operation/commit/sequence_snapshot_properties.h" #include "paimon/core/operation/expire_snapshots.h" @@ -80,6 +81,7 @@ #include "paimon/core/utils/duration.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/file_store_write.h" #include "paimon/fs/file_system.h" #include "paimon/logging.h" #include "paimon/metrics.h" @@ -383,7 +385,8 @@ Result FileStoreCommitImpl::FilterAndCommit( std::optional watermark) { std::vector> committables; for (const auto& [identifier, msgs] : commit_identifier_and_messages) { - committables.push_back(CreateManifestCommittable(identifier, msgs, watermark)); + committables.push_back( + CreateManifestCommittable(identifier, msgs, watermark, /*properties=*/{})); } std::vector> sorted_committables = committables; @@ -535,7 +538,7 @@ Status FileStoreCommitImpl::Overwrite( const std::vector>& commit_messages, int64_t identifier, std::optional watermark) { std::shared_ptr committable = - CreateManifestCommittable(identifier, commit_messages, watermark); + CreateManifestCommittable(identifier, commit_messages, watermark, /*properties=*/{}); PAIMON_LOG_INFO(logger_, "Ready to overwrite to table %s, number of commit messages: %zu", table_name_.c_str(), committable->FileCommittables().size()); PAIMON_ASSIGN_OR_RAISE(std::string committable_str, committable->ToString()); @@ -576,7 +579,7 @@ Result FileStoreCommitImpl::FilterAndOverwrite( const std::vector>& commit_messages, int64_t identifier, std::optional watermark) { std::shared_ptr committable = - CreateManifestCommittable(identifier, commit_messages, watermark); + CreateManifestCommittable(identifier, commit_messages, watermark, /*properties=*/{}); PAIMON_LOG_INFO(logger_, "Ready to overwrite to table %s, number of commit messages: %zu", table_name_.c_str(), committable->FileCommittables().size()); PAIMON_ASSIGN_OR_RAISE(std::string committable_str, committable->ToString()); @@ -872,7 +875,51 @@ Status FileStoreCommitImpl::Commit( const std::vector>& commit_messages, int64_t identifier, std::optional watermark) { std::shared_ptr committable = - CreateManifestCommittable(identifier, commit_messages, watermark); + CreateManifestCommittable(identifier, commit_messages, watermark, /*properties=*/{}); + return Commit(committable, /*check_append_files=*/false); +} + +Status FileStoreCommitImpl::CommitWithProgress( + const std::vector& realtime_commits, int64_t identifier, + std::optional watermark) { + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager_->LatestSnapshot()); + PAIMON_ASSIGN_OR_RAISE(RealtimeSnapshotProperties::OffsetMap committed_offsets, + RealtimeSnapshotProperties::ReadOffsets(latest_snapshot, fs_)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeSnapshotProperties::ValidatedCommitProgress validated_progress, + RealtimeSnapshotProperties::ValidateProgress(realtime_commits, committed_offsets)); + + std::vector> commit_messages; + commit_messages.reserve(validated_progress.ordered_commits.size()); + for (const RealtimeCommitProgress& realtime_commit : validated_progress.ordered_commits) { + if (!realtime_commit.commit_message) { + return Status::Invalid("real-time commit message is null"); + } + auto commit_message = + std::dynamic_pointer_cast(realtime_commit.commit_message); + if (!commit_message) { + return Status::Invalid("fail to cast real-time commit message to impl"); + } + PAIMON_ASSIGN_OR_RAISE(std::string message_partition, + path_factory_->GetPartitionString(commit_message->Partition())); + message_partition = PartitionBucket::NormalizePartition(std::move(message_partition)); + if (message_partition != realtime_commit.partition || + commit_message->Bucket() != realtime_commit.bucket) { + return Status::Invalid( + "real-time commit progress does not match commit message partition and bucket"); + } + commit_messages.push_back(realtime_commit.commit_message); + } + + std::map properties; + if (!validated_progress.delta_offsets.empty()) { + PAIMON_ASSIGN_OR_RAISE( + properties[RealtimeSnapshotProperties::kOffsetsDeltaKey], + RealtimeSnapshotProperties::SerializeOffsets(validated_progress.delta_offsets)); + } + std::shared_ptr committable = + CreateManifestCommittable(identifier, commit_messages, watermark, properties); return Commit(committable, /*check_append_files=*/false); } @@ -1193,7 +1240,12 @@ Result FileStoreCommitImpl::TryCommitOnce( statistics = std::nullopt; } - std::map snapshot_properties = properties; + using SnapshotProperties = std::map; + PAIMON_ASSIGN_OR_RAISE( + SnapshotProperties snapshot_properties, + RealtimeSnapshotProperties::MergeOffsets( + properties, latest_snapshot, fs_, + RealtimeSnapshotProperties::OffsetsDirectory(root_path_, snapshot_manager_->Branch()))); if (options_.WriteSequenceNumberInitMode() == CoreOptions::SequenceNumberInitMode::SNAPSHOT) { PAIMON_ASSIGN_OR_RAISE(std::optional latest_max_sequence_number, SequenceSnapshotProperties::MaxSequenceNumber(latest_snapshot)); @@ -1316,12 +1368,9 @@ void FileStoreCommitImpl::CleanUpTmpManifests( std::shared_ptr FileStoreCommitImpl::CreateManifestCommittable( int64_t identifier, const std::vector>& commit_messages, - std::optional watermark) { - auto committable = std::make_shared(identifier, watermark); - for (const auto& commit_message : commit_messages) { - committable->AddFileCommittable(commit_message); - } - return committable; + std::optional watermark, const std::map& properties) { + return std::make_shared(identifier, watermark, properties, + commit_messages); } Result FileStoreCommitImpl::CollectChanges( diff --git a/src/paimon/core/operation/file_store_commit_impl.h b/src/paimon/core/operation/file_store_commit_impl.h index c0c93c9a..246455c8 100644 --- a/src/paimon/core/operation/file_store_commit_impl.h +++ b/src/paimon/core/operation/file_store_commit_impl.h @@ -106,6 +106,9 @@ class FileStoreCommitImpl : public FileStoreCommit { int64_t commit_identifier, std::optional watermark = std::nullopt) override; + Status CommitWithProgress(const std::vector& realtime_commits, + int64_t commit_identifier, std::optional watermark) override; + Result FilterAndCommit( const std::map>>& commit_identifier_and_messages, @@ -172,7 +175,7 @@ class FileStoreCommitImpl : public FileStoreCommit { std::shared_ptr CreateManifestCommittable( int64_t identifier, const std::vector>& commit_messages, - std::optional watermark); + std::optional watermark, const std::map& properties); Result CollectChanges( const std::vector>& commit_messages); diff --git a/src/paimon/core/operation/file_store_commit_impl_test.cpp b/src/paimon/core/operation/file_store_commit_impl_test.cpp index 405b8280..2186dc88 100644 --- a/src/paimon/core/operation/file_store_commit_impl_test.cpp +++ b/src/paimon/core/operation/file_store_commit_impl_test.cpp @@ -1695,7 +1695,8 @@ TEST_F(FileStoreCommitImplTest, TestCreateManifestCommittable) { auto commit_impl = std::dynamic_pointer_cast( std::shared_ptr(std::move(commit))); std::vector> msgs; - auto committable = commit_impl->CreateManifestCommittable(1, msgs, 30); + const std::map properties; + auto committable = commit_impl->CreateManifestCommittable(1, msgs, 30, properties); ASSERT_TRUE(committable); EXPECT_EQ(1, committable->Identifier()); EXPECT_EQ(30, committable->Watermark().value()); @@ -1763,7 +1764,8 @@ TEST_F(FileStoreCommitImplTest, TestFilterCommitted) { "/orc/append_09.db/append_09/commit_messages/" "commit_messages-01", /*version=*/3); - auto committable = commit_impl->CreateManifestCommittable(1, msgs, std::nullopt); + const std::map properties; + auto committable = commit_impl->CreateManifestCommittable(1, msgs, std::nullopt, properties); std::vector> committables = {committable}; // Test with no previous snapshots @@ -1793,14 +1795,15 @@ TEST_F(FileStoreCommitImplTest, TestFilterCommittedWithMultipleCommittables) { "/orc/append_09.db/append_09/commit_messages/" "commit_messages-01", /*version=*/3); - auto committable1 = commit_impl->CreateManifestCommittable(1, msgs1, std::nullopt); + const std::map properties; + auto committable1 = commit_impl->CreateManifestCommittable(1, msgs1, std::nullopt, properties); std::vector> msgs2 = GetCommitMessages(paimon::test::GetDataDir() + "/orc/append_09.db/append_09/commit_messages/" "commit_messages-02", /*version=*/3); - auto committable2 = commit_impl->CreateManifestCommittable(2, msgs2, std::nullopt); + auto committable2 = commit_impl->CreateManifestCommittable(2, msgs2, std::nullopt, properties); std::vector> committables = {committable1, committable2}; @@ -1838,8 +1841,9 @@ TEST_F(FileStoreCommitImplTest, TestFilterCommittedRejectsDuplicateIdentifiers) "commit_messages-02", /*version=*/3); - auto committable1 = commit_impl->CreateManifestCommittable(1, msgs1, std::nullopt); - auto committable2 = commit_impl->CreateManifestCommittable(1, msgs2, std::nullopt); + const std::map properties; + auto committable1 = commit_impl->CreateManifestCommittable(1, msgs1, std::nullopt, properties); + auto committable2 = commit_impl->CreateManifestCommittable(1, msgs2, std::nullopt, properties); std::vector> committables = {committable1, committable2}; ASSERT_NOK_WITH_MSG(commit_impl->FilterCommitted(committables), diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 5b091ebf..6e79a3f2 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -32,6 +32,7 @@ #include "paimon/core/mergetree/compact/merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/append_only_file_store_write.h" +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" @@ -55,6 +56,10 @@ struct KeyValue; template class MergeFunctionWrapper; +Result> FileStoreWrite::PrepareCommitWithProgress(int64_t) { + return Status::Invalid("prepare commit with progress requires a real-time writer"); +} + Result> FileStoreWrite::Create(std::unique_ptr ctx) { if (ctx == nullptr) { return Status::Invalid("write context is null pointer"); @@ -103,6 +108,14 @@ Result> FileStoreWrite::Create(std::unique_ptrGetMemoryPool())); auto snapshot_manager = std::make_shared(options.GetFileSystem(), ctx->GetRootPath(), branch); + RealtimeSnapshotProperties::OffsetMap realtime_committed_offsets; + if (ctx->GetRealtimeContext()) { + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager->LatestSnapshot()); + PAIMON_ASSIGN_OR_RAISE( + realtime_committed_offsets, + RealtimeSnapshotProperties::ReadOffsets(latest_snapshot, options.GetFileSystem())); + } std::shared_ptr io_manager; const auto& io_temp_dir = ctx->GetTempDirectory(); @@ -121,6 +134,17 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { + if (options.GetBucket() <= 0) { + return Status::Invalid("real-time append write requires fixed bucket mode"); + } + if (options.DataEvolutionEnabled() || !ctx->GetWriteSchema().empty()) { + return Status::Invalid("real-time append write does not support data evolution"); + } + if (options.DeletionVectorsEnabled()) { + return Status::Invalid("real-time append write does not support deletion vectors"); + } + } std::shared_ptr write_schema = arrow_schema; const auto& write_field_names = ctx->GetWriteSchema(); if (!write_field_names.empty()) { @@ -157,10 +181,13 @@ Result> FileStoreWrite::Create(std::unique_ptrGetCommitUser(), ctx->GetRootPath(), schema, arrow_schema, write_schema, partition_schema, dv_maintainer_factory, io_manager, options, ignore_previous_files, - ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), ctx->GetExecutor(), - ctx->GetMemoryPool()); + ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), ctx->GetRealtimeContext(), + realtime_committed_offsets, ctx->GetExecutor(), ctx->GetMemoryPool()); } else { // pk table + if (ctx->GetRealtimeContext()) { + return Status::Invalid("real-time write currently supports append tables only"); + } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( snapshot_manager, schema_manager, ctx->GetCommitUser(), ctx->GetRootPath(), schema, diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index e8886b42..4e31f5a8 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -136,6 +136,9 @@ std::optional InternalReadContext::TryResolveSpecialFieldById( } return std::nullopt; } + if (field_id == SpecialFields::Offset().Id()) { + return SpecialFields::Offset(); + } if (field_id == SpecialFields::IndexScore().Id()) { if (core_options.DataEvolutionEnabled()) { return SpecialFields::IndexScore(); @@ -162,6 +165,9 @@ std::optional InternalReadContext::TryResolveSpecialFieldByName( } return std::nullopt; } + if (name == SpecialFields::Offset().Name()) { + return SpecialFields::Offset(); + } if (name == SpecialFields::IndexScore().Name()) { if (core_options.DataEvolutionEnabled()) { return SpecialFields::IndexScore(); diff --git a/src/paimon/core/operation/write_context.cpp b/src/paimon/core/operation/write_context.cpp index aab37d32..457f8d5b 100644 --- a/src/paimon/core/operation/write_context.cpp +++ b/src/paimon/core/operation/write_context.cpp @@ -41,6 +41,7 @@ WriteContext::WriteContext(const std::string& root_path, const std::string& comm const std::string& temp_directory, const std::shared_ptr& specific_file_system, const std::map& fs_scheme_to_identifier_map, + const std::shared_ptr& realtime_context, const std::map& options) : root_path_(root_path), commit_user_(commit_user), @@ -56,6 +57,7 @@ WriteContext::WriteContext(const std::string& root_path, const std::string& comm temp_directory_(temp_directory), specific_file_system_(specific_file_system), fs_scheme_to_identifier_map_(fs_scheme_to_identifier_map), + realtime_context_(realtime_context), options_(options) {} WriteContext::~WriteContext() = default; @@ -77,6 +79,7 @@ class WriteContextBuilder::Impl { write_schema_.clear(); fs_scheme_to_identifier_map_.clear(); specific_file_system_.reset(); + realtime_context_.reset(); options_.clear(); } @@ -95,6 +98,7 @@ class WriteContextBuilder::Impl { std::string temp_directory_; std::map fs_scheme_to_identifier_map_; std::shared_ptr specific_file_system_; + std::shared_ptr realtime_context_; std::map options_; }; @@ -183,6 +187,12 @@ WriteContextBuilder& WriteContextBuilder::WithFileSystem( return *this; } +WriteContextBuilder& WriteContextBuilder::WithRealtimeContext( + const std::shared_ptr& realtime_context) { + impl_->realtime_context_ = realtime_context; + return *this; +} + Result> WriteContextBuilder::Finish() { PAIMON_ASSIGN_OR_RAISE(impl_->root_path_, PathUtil::NormalizePath(impl_->root_path_)); if (impl_->root_path_.empty()) { @@ -198,7 +208,7 @@ Result> WriteContextBuilder::Finish() { impl_->ignore_num_bucket_check_, impl_->ignore_previous_files_, enable_multi_thread_spill, impl_->write_id_, impl_->branch_, impl_->write_schema_, impl_->memory_pool_, impl_->executor_, impl_->temp_directory_, impl_->specific_file_system_, - impl_->fs_scheme_to_identifier_map_, impl_->options_); + impl_->fs_scheme_to_identifier_map_, impl_->realtime_context_, impl_->options_); impl_->Reset(); return ctx; } diff --git a/src/paimon/core/realtime/arrow_mem_indexer.cpp b/src/paimon/core/realtime/arrow_mem_indexer.cpp new file mode 100644 index 00000000..136f0e2e --- /dev/null +++ b/src/paimon/core/realtime/arrow_mem_indexer.cpp @@ -0,0 +1,204 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/arrow_mem_indexer.h" + +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/macros.h" + +namespace paimon { +namespace { + +uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t result = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + result += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + result += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + result += GetArrayMemoryUsage(data->dictionary); + } + return result; +} + +} // namespace + +class ArrowMemIndexer::Segment : public RealtimeSegmentHandle { + public: + Segment(const Range& offset_range, std::vector&& batches) + : offset_range_(offset_range), batches_(std::move(batches)) {} + + Range GetOffsetRange() const override { + return offset_range_; + } + + const std::vector& GetBatches() const { + return batches_; + } + + private: + Range offset_range_; + std::vector batches_; +}; + +class ArrowMemIndexer::CommitBatchReader : public BatchReader { + public: + CommitBatchReader(const std::shared_ptr& segment, + const std::shared_ptr& arrow_pool) + : segment_(segment), arrow_pool_(arrow_pool), metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!segment_ || next_batch_ >= segment_->GetBatches().size()) { + return MakeEofBatch(); + } + const StoredBatch& stored = segment_->GetBatches()[next_batch_++]; + int64_t row_count = stored.data->length(); + + arrow::Int8Builder row_kind_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Reserve(row_count)); + if (stored.row_kinds.empty()) { + for (int64_t i = 0; i < row_count; ++i) { + row_kind_builder.UnsafeAppend(static_cast(RecordBatch::RowKind::INSERT)); + } + } else { + for (RecordBatch::RowKind row_kind : stored.row_kinds) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + row_kind_builder.Append(static_cast(row_kind))); + } + } + std::shared_ptr row_kind_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array)); + + arrow::ArrayVector 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, stored.data->null_bitmap(), + stored.data->null_count(), stored.data->offset())); + 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())); + return ReadBatch(std::move(c_array), std::move(c_schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + + void Close() override { + segment_.reset(); + } + + private: + std::shared_ptr segment_; + std::shared_ptr arrow_pool_; + std::shared_ptr metrics_; + size_t next_batch_ = 0; +}; + +ArrowMemIndexer::ArrowMemIndexer(const std::shared_ptr& write_schema, + const std::shared_ptr& arrow_pool) + : write_schema_(write_schema), arrow_pool_(arrow_pool) {} + +Status ArrowMemIndexer::Write(RealtimeWriteBatch&& write_batch) { + if (!write_batch.batch) { + return Status::Invalid("real-time write batch is null"); + } + int64_t row_count = write_batch.batch->GetData()->length; + if (write_batch.offset_range.Count() != row_count) { + return Status::Invalid("real-time offset range 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()))); + std::shared_ptr struct_array = + std::dynamic_pointer_cast(data); + if (!struct_array) { + return Status::Invalid("real-time write data is not a StructArray"); + } + + if (closed_) { + return Status::Invalid("mem indexer is closed"); + } + if (building_range_ && write_batch.offset_range.from != building_range_->to + 1) { + return Status::Invalid("real-time offset ranges must be contiguous"); + } + building_memory_usage_ += GetArrayMemoryUsage(struct_array->data()); + building_batches_.push_back( + StoredBatch{std::move(struct_array), write_batch.batch->GetRowKind()}); + if (!building_range_) { + building_range_ = write_batch.offset_range; + } else { + building_range_ = Range(building_range_->from, write_batch.offset_range.to); + } + return Status::OK(); +} + +Result>> ArrowMemIndexer::SealForCommit() { + if (closed_) { + return Status::Invalid("mem indexer is closed"); + } + if (building_batches_.empty()) { + return std::optional>(); + } + auto segment = std::make_shared(building_range_.value(), std::move(building_batches_)); + building_batches_.clear(); + building_range_.reset(); + building_memory_usage_ = 0; + return std::optional>(std::move(segment)); +} + +Result>> ArrowMemIndexer::CreateCommitReaders( + const std::shared_ptr& segment) { + std::shared_ptr arrow_segment = std::dynamic_pointer_cast(segment); + if (!arrow_segment) { + return Status::Invalid("segment was not created by the Arrow mem indexer"); + } + std::vector> readers; + readers.push_back(std::make_unique(arrow_segment, arrow_pool_)); + return readers; +} + +uint64_t ArrowMemIndexer::GetMemoryUsage() const { + return building_memory_usage_; +} + +Status ArrowMemIndexer::Close() { + building_batches_.clear(); + building_memory_usage_ = 0; + building_range_.reset(); + closed_ = true; + return Status::OK(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/arrow_mem_indexer.h b/src/paimon/core/realtime/arrow_mem_indexer.h new file mode 100644 index 00000000..0663cab9 --- /dev/null +++ b/src/paimon/core/realtime/arrow_mem_indexer.h @@ -0,0 +1,68 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/realtime/mem_indexer.h" + +namespace arrow { +class MemoryPool; +class Schema; +class StructArray; +} // namespace arrow + +namespace paimon { + +/// Internal Arrow-backed implementation of the default `MemIndexer`. +class ArrowMemIndexer : public MemIndexer { + public: + ArrowMemIndexer(const std::shared_ptr& write_schema, + const std::shared_ptr& arrow_pool); + + Status Write(RealtimeWriteBatch&& write_batch) override; + + Result>> SealForCommit() override; + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + + uint64_t GetMemoryUsage() const override; + + Status Close() override; + + private: + struct StoredBatch { + std::shared_ptr data; + std::vector row_kinds; + }; + + class Segment; + class CommitBatchReader; + + std::shared_ptr write_schema_; + std::shared_ptr arrow_pool_; + std::vector building_batches_; + std::optional building_range_; + uint64_t building_memory_usage_ = 0; + bool closed_ = false; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp b/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp new file mode 100644 index 00000000..faa09a57 --- /dev/null +++ b/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp @@ -0,0 +1,42 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/realtime/arrow_mem_indexer_factory.h" + +#include "arrow/c/bridge.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/core/realtime/arrow_mem_indexer.h" +#include "paimon/macros.h" + +namespace paimon { + +Result> ArrowMemIndexerFactory::Create( + ArrowSchema* write_schema, const std::map&, + const std::shared_ptr& memory_pool) { + if (!write_schema || !write_schema->release) { + return Status::Invalid("mem indexer write schema is null"); + } + if (!memory_pool) { + return Status::Invalid("mem indexer memory pool is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, + arrow::ImportSchema(write_schema)); + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + return std::make_shared(imported_schema, arrow_pool); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/partition_bucket.h b/src/paimon/core/realtime/partition_bucket.h new file mode 100644 index 00000000..be04f9b6 --- /dev/null +++ b/src/paimon/core/realtime/partition_bucket.h @@ -0,0 +1,52 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/common/utils/path_util.h" + +namespace paimon { + +/// Identifies one partition-bucket after normalizing the partition path. +struct PartitionBucket { + PartitionBucket(std::string partition, int32_t bucket) + : partition(NormalizePartition(std::move(partition))), bucket(bucket) {} + + static std::string NormalizePartition(std::string partition) { + PathUtil::TrimLastDelim(&partition); + return partition; + } + + bool operator<(const PartitionBucket& other) const { + if (partition != other.partition) { + return partition < other.partition; + } + return bucket < other.bucket; + } + + bool operator==(const PartitionBucket& other) const { + return partition == other.partition && bucket == other.bucket; + } + + std::string partition; + int32_t bucket = -1; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp new file mode 100644 index 00000000..4226307c --- /dev/null +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -0,0 +1,231 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/realtime_append_only_writer.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/append/append_only_writer.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" +#include "paimon/realtime/realtime_context.h" + +namespace paimon { + +Result> RealtimeAppendOnlyWriter::Create( + const std::string& 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, + const std::shared_ptr& realtime_write_schema, + const std::map& options, int64_t next_offset, + const std::shared_ptr& memory_pool) { + if (!realtime_context) { + return Status::Invalid("real-time context is null"); + } + if (next_offset < 0) { + return Status::Invalid("next real-time offset must be non-negative"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr mem_indexer, + realtime_context->GetOrCreateMemIndexer( + partition, bucket, std::move(write_schema), options, memory_pool)); + return std::shared_ptr(new RealtimeAppendOnlyWriter( + mem_indexer, file_writer, input_schema, realtime_write_schema, next_offset, memory_pool)); +} + +RealtimeAppendOnlyWriter::RealtimeAppendOnlyWriter( + const std::shared_ptr& mem_indexer, + 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& memory_pool) + : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + mem_indexer_(mem_indexer), + file_writer_(file_writer), + input_schema_(input_schema), + realtime_write_schema_(realtime_write_schema), + next_offset_(next_offset) {} + +Status RealtimeAppendOnlyWriter::Write(std::unique_ptr&& batch) { + for (RecordBatch::RowKind row_kind : batch->GetRowKind()) { + if (row_kind != RecordBatch::RowKind::INSERT) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* kind, + RowKind::FromByteValue(static_cast(row_kind))); + return Status::Invalid("Append only writer can not accept record batch with RowKind ", + kind->Name()); + } + } + + int64_t row_count = batch->GetData()->length; + if (row_count == 0) { + return Status::OK(); + } + int64_t first_offset = next_offset_.fetch_add(row_count); + Range range(first_offset, first_offset + row_count - 1); + std::lock_guard lock(mem_indexer_mutex_); + return mem_indexer_->Write(RealtimeWriteBatch{std::move(batch), range}); +} + +Result RealtimeAppendOnlyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard lock(prepare_mutex_); + std::optional> segment; + { + std::lock_guard mem_indexer_lock(mem_indexer_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> sealed_segment, + mem_indexer_->SealForCommit()); + segment = std::move(sealed_segment); + } + if (segment) { + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value())); + } + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, file_writer_->PrepareCommit(wait_compaction)); + if (segment) { + increment.SetRealtimeOffsetRange(segment.value()->GetOffsetRange()); + } + return increment; +} + +Status RealtimeAppendOnlyWriter::FlushSegment( + const std::shared_ptr& segment) { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + mem_indexer_->CreateCommitReaders(segment)); + ConcatBatchReader reader(std::move(readers), memory_pool_); + ScopeGuard reader_guard([&reader]() { reader.Close(); }); + const Range offset_range = segment->GetOffsetRange(); + int64_t emitted_rows = 0; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader.NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, + arrow::ImportArray(c_array.get(), c_schema.get())); + std::shared_ptr struct_array = + std::dynamic_pointer_cast(imported); + if (!struct_array) { + return Status::Invalid("mem indexer commit reader returned a non-StructArray"); + } + std::shared_ptr value_kind = + struct_array->GetFieldByName(SpecialFields::ValueKind().Name()); + if (!value_kind || value_kind->type_id() != arrow::Type::INT8) { + return Status::Invalid( + "mem indexer commit reader must return an INT8 _VALUE_KIND field"); + } + std::shared_ptr row_kinds = + std::static_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 mem indexer commit reader returned a non-INSERT row"); + } + } + PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( + struct_array, SpecialFields::ValueKind().Name())); + if (!struct_array->type()->Equals(arrow::struct_(input_schema_->fields()))) { + return Status::Invalid( + "mem indexer commit reader schema does not match table write schema"); + } + + int64_t row_count = struct_array->length(); + if (row_count > offset_range.Count() - emitted_rows) { + return Status::Invalid( + "mem indexer commit readers returned more rows than the sealed offset range"); + } + arrow::Int64Builder offset_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offset_builder.Reserve(row_count)); + for (int64_t i = 0; i < row_count; ++i) { + offset_builder.UnsafeAppend(offset_range.from + emitted_rows + i); + } + std::shared_ptr offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(offset_builder.Finish(&offset_array)); + + arrow::ArrayVector arrays = struct_array->fields(); + arrays.insert(arrays.begin(), offset_array); + arrow::FieldVector fields = struct_array->struct_type()->fields(); + fields.insert(fields.begin(), + DataField::ConvertDataFieldToArrowField(SpecialFields::Offset())); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + struct_array, + arrow::StructArray::Make(arrays, fields, struct_array->null_bitmap(), + struct_array->null_count(), struct_array->offset())); + if (!struct_array->type()->Equals(arrow::struct_(realtime_write_schema_->fields()))) { + return Status::Invalid("failed to build the real-time data-file write schema"); + } + emitted_rows += row_count; + + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, output.get())); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr record_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(file_writer_->Write(std::move(record_batch))); + } + if (emitted_rows != offset_range.Count()) { + return Status::Invalid( + "mem indexer commit readers returned fewer rows than the sealed offset range"); + } + return Status::OK(); +} + +Status RealtimeAppendOnlyWriter::Compact(bool) { + return Status::Invalid("real-time append write does not support explicit compaction"); +} + +uint64_t RealtimeAppendOnlyWriter::GetMemoryUsage() const { + // The first implementation does not spill sealed or building segments through + // WriterMemoryManager. + return 0; +} + +Status RealtimeAppendOnlyWriter::FlushMemory() { + return Status::OK(); +} + +Result RealtimeAppendOnlyWriter::CompactNotCompleted() { + return file_writer_->CompactNotCompleted(); +} + +Status RealtimeAppendOnlyWriter::Sync() { + return file_writer_->Sync(); +} + +Status RealtimeAppendOnlyWriter::Close() { + { + std::lock_guard lock(mem_indexer_mutex_); + PAIMON_RETURN_NOT_OK(mem_indexer_->Close()); + } + return file_writer_->Close(); +} + +std::shared_ptr RealtimeAppendOnlyWriter::GetMetrics() const { + return file_writer_->GetMetrics(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h new file mode 100644 index 00000000..2560a049 --- /dev/null +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -0,0 +1,90 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/mem_indexer.h" + +struct ArrowSchema; + +namespace arrow { +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { + +class AppendOnlyWriter; +class MemoryPool; +class RealtimeContext; + +class RealtimeAppendOnlyWriter : public BatchWriter { + public: + static Result> Create( + const std::string& 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, + const std::shared_ptr& realtime_write_schema, + const std::map& options, int64_t next_offset, + const std::shared_ptr& memory_pool); + + Status Write(std::unique_ptr&& batch) override; + + Result PrepareCommit(bool wait_compaction) override; + + Status Compact(bool full_compaction) override; + + uint64_t GetMemoryUsage() const override; + + Status FlushMemory() override; + + Result CompactNotCompleted() override; + + Status Sync() override; + + Status Close() override; + + std::shared_ptr GetMetrics() const override; + + private: + RealtimeAppendOnlyWriter(const std::shared_ptr& mem_indexer, + 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& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment); + + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr mem_indexer_; + std::shared_ptr file_writer_; + std::shared_ptr input_schema_; + std::shared_ptr realtime_write_schema_; + std::atomic next_offset_; + std::mutex mem_indexer_mutex_; + std::mutex prepare_mutex_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_context.cpp b/src/paimon/core/realtime/realtime_context.cpp new file mode 100644 index 00000000..9a9a0dcd --- /dev/null +++ b/src/paimon/core/realtime/realtime_context.cpp @@ -0,0 +1,89 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/realtime/realtime_context.h" + +#include +#include +#include + +#include "paimon/arrow/abi.h" +#include "paimon/core/realtime/partition_bucket.h" +#include "paimon/macros.h" +#include "paimon/realtime/arrow_mem_indexer_factory.h" +#include "paimon/realtime/mem_indexer.h" +#include "paimon/status.h" + +namespace paimon { + +class RealtimeContext::Impl { + public: + explicit Impl(const std::shared_ptr& factory) : factory_(factory) {} + + Result> GetOrCreateMemIndexer( + const std::string& partition, int32_t bucket, std::unique_ptr write_schema, + const std::map& options, + const std::shared_ptr& memory_pool) { + std::lock_guard lock(mutex_); + const PartitionBucket key(partition, bucket); + auto iter = indexers_.find(key); + if (iter != indexers_.end()) { + if (write_schema && write_schema->release) { + write_schema->release(write_schema.get()); + } + return iter->second; + } + Result> indexer_result = + factory_->Create(write_schema.get(), options, memory_pool); + if (write_schema && write_schema->release) { + write_schema->release(write_schema.get()); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr indexer, std::move(indexer_result)); + indexers_.emplace(key, indexer); + return indexer; + } + + private: + std::shared_ptr factory_; + std::mutex mutex_; + std::map> indexers_; +}; + +Result> RealtimeContext::Create() { + return Create(std::make_shared()); +} + +Result> RealtimeContext::Create( + const std::shared_ptr& factory) { + if (!factory) { + return Status::Invalid("mem indexer factory is null"); + } + return std::shared_ptr(new RealtimeContext(std::make_unique(factory))); +} + +RealtimeContext::RealtimeContext(std::unique_ptr&& impl) : impl_(std::move(impl)) {} + +RealtimeContext::~RealtimeContext() = default; + +Result> RealtimeContext::GetOrCreateMemIndexer( + const std::string& partition, int32_t bucket, std::unique_ptr write_schema, + const std::map& options, + const std::shared_ptr& memory_pool) { + return impl_->GetOrCreateMemIndexer(partition, bucket, std::move(write_schema), options, + memory_pool); +} + +} // namespace paimon diff --git a/src/paimon/core/utils/commit_increment.h b/src/paimon/core/utils/commit_increment.h index e3718218..7fb73b93 100644 --- a/src/paimon/core/utils/commit_increment.h +++ b/src/paimon/core/utils/commit_increment.h @@ -20,10 +20,12 @@ #pragma once #include +#include #include "paimon/core/compact/compact_deletion_file.h" #include "paimon/core/io/compact_increment.h" #include "paimon/core/io/data_increment.h" +#include "paimon/utils/range.h" namespace paimon { @@ -56,10 +58,19 @@ class CommitIncrement { return compact_deletion_file_; } + const std::optional& GetRealtimeOffsetRange() const { + return realtime_offset_range_; + } + + void SetRealtimeOffsetRange(const Range& offset_range) { + realtime_offset_range_ = offset_range; + } + private: DataIncrement data_increment_; CompactIncrement compact_increment_; std::shared_ptr compact_deletion_file_; + std::optional realtime_offset_range_; }; } // namespace paimon diff --git a/test/inte/CMakeLists.txt b/test/inte/CMakeLists.txt index 85fb8f30..76e57149 100644 --- a/test/inte/CMakeLists.txt +++ b/test/inte/CMakeLists.txt @@ -36,6 +36,13 @@ if(PAIMON_BUILD_TESTS) test_utils_static ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(realtime_write_inte_test + STATIC_LINK_LIBS + paimon_shared + ${TEST_STATIC_LINK_LIBS} + test_utils_static + ${GTEST_LINK_TOOLCHAIN}) + add_paimon_test(global_index_test STATIC_LINK_LIBS paimon_shared diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp new file mode 100644 index 00000000..8e5736e7 --- /dev/null +++ b/test/inte/realtime_write_inte_test.cpp @@ -0,0 +1,423 @@ +/* + * Copyright 2026-present Alibaba Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/catalog/catalog.h" +#include "paimon/catalog/identifier.h" +#include "paimon/commit_context.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/core/core_options.h" +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" +#include "paimon/core/utils/snapshot_manager.h" +#include "paimon/defs.h" +#include "paimon/file_store_commit.h" +#include "paimon/file_store_write.h" +#include "paimon/read_context.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/record_batch.h" +#include "paimon/scan_context.h" +#include "paimon/table/source/table_read.h" +#include "paimon/table/source/table_scan.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/write_context.h" + +namespace paimon::test { + +class RealtimeWriteInteTest : public ::testing::Test { + protected: + using Row = std::tuple; + using ReadRow = std::tuple; + + void SetUp() override { + dir_ = UniqueTestDirectory::Create("local"); + ASSERT_NE(nullptr, dir_); + table_path_ = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + fields_ = {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), + arrow::field("pt", arrow::utf8())}; + schema_ = arrow::schema(fields_); + options_ = { + {Options::MANIFEST_FORMAT, "orc"}, {Options::FILE_FORMAT, "orc"}, + {Options::FILE_SYSTEM, "local"}, {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, {Options::TARGET_FILE_SIZE, "1048576"}, + }; + } + + void TearDown() override { + dir_.reset(); + } + + void CreateTable(const std::vector& partition_keys) const { + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, + Catalog::Create(dir_->Str(), options_)); + ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false)); + ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), &c_schema, partition_keys, + /*primary_keys=*/{}, options_, + /*ignore_if_exists=*/false)); + } + + Result> CreateRealtimeWriter() const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path_, commit_user_); + builder.SetOptions(options_).WithStreamingMode(true).WithRealtimeContext(realtime_context); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + return FileStoreWrite::Create(std::move(context)); + } + + Result> MakeBatch(const std::vector& rows, + bool partitioned) const { + if (rows.empty()) { + return Status::Invalid("cannot create an empty test batch"); + } + const std::string& partition = std::get<2>(rows.front()); + std::string json = "["; + for (size_t i = 0; i < rows.size(); ++i) { + const auto& [id, payload, pt] = rows[i]; + if (pt != partition) { + return Status::Invalid("one test batch must contain only one partition"); + } + if (i > 0) { + json += ","; + } + json += "[" + std::to_string(id) + ",\"" + payload + "\",\"" + pt + "\"]"; + } + json += "]"; + + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + if (partitioned) { + builder.SetPartition({{"pt", partition}}); + } + return builder.SetBucket(0).Finish(); + } + + static std::vector MakeRows(int64_t first_id, int64_t count, + const std::string& partition) { + std::vector rows; + rows.reserve(count); + for (int64_t i = 0; i < count; ++i) { + int64_t id = first_id + i; + rows.emplace_back(id, "value-" + std::to_string(id), partition); + } + return rows; + } + + Status Commit(const std::vector& realtime_commits, + int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, commit_user_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + return commit->CommitWithProgress(realtime_commits, commit_identifier, + /*watermark=*/std::nullopt); + } + + static std::vector WithExpectedOffsets(const std::vector& rows) { + std::map next_offsets; + std::vector result; + result.reserve(rows.size()); + for (const Row& row : rows) { + const auto& [id, payload, partition] = row; + int64_t offset = next_offsets[partition]++; + result.emplace_back(offset, id, payload, partition); + } + return result; + } + + Result> ReadRows() const { + ScanContextBuilder scan_builder(table_path_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, + scan_builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan, + TableScan::Create(std::move(scan_context))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, scan->CreatePlan()); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_).SetReadFieldNames({"_OFFSET", "id", "payload", "pt"}); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + table_read->CreateReader(plan->Splits())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, + ReadResultCollector::CollectResult(reader.get())); + reader->Close(); + + std::vector rows; + if (!result) { + return rows; + } + for (const std::shared_ptr& chunk : result->chunks()) { + std::shared_ptr data = + std::dynamic_pointer_cast(chunk); + if (!data || data->num_fields() != 5) { + return Status::Invalid("unexpected real-time test read schema"); + } + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(data->field(0)); + std::shared_ptr offsets = + std::dynamic_pointer_cast(data->field(1)); + std::shared_ptr ids = + std::dynamic_pointer_cast(data->field(2)); + std::shared_ptr payloads = + std::dynamic_pointer_cast(data->field(3)); + std::shared_ptr partitions = + std::dynamic_pointer_cast(data->field(4)); + if (!row_kinds || !offsets || !ids || !payloads || !partitions) { + return Status::Invalid("unexpected real-time test read field type"); + } + for (int64_t i = 0; i < data->length(); ++i) { + if (row_kinds->IsNull(i) || + row_kinds->Value(i) != static_cast(RecordBatch::RowKind::INSERT) || + offsets->IsNull(i) || ids->IsNull(i) || payloads->IsNull(i) || + partitions->IsNull(i)) { + return Status::Invalid("unexpected null or row kind in real-time test result"); + } + rows.emplace_back(offsets->Value(i), ids->Value(i), payloads->GetString(i), + partitions->GetString(i)); + } + } + return rows; + } + + Result ReadCommittedOffsets() const { + PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional snapshot, snapshot_manager.LatestSnapshot()); + return RealtimeSnapshotProperties::ReadOffsets(snapshot, options.GetFileSystem()); + } + + void FinalizeCommitAndCheck(FileStoreWrite* writer, + std::vector realtime_commits, + int64_t prepare_identifier, std::vector expected_rows) const { + ASSERT_OK_AND_ASSIGN(std::vector final_commits, + writer->PrepareCommitWithProgress(prepare_identifier)); + realtime_commits.insert(realtime_commits.end(), + std::make_move_iterator(final_commits.begin()), + std::make_move_iterator(final_commits.end())); + // Verify that commit orders prepared offset ranges instead of relying on caller order. + std::reverse(realtime_commits.begin(), realtime_commits.end()); + ASSERT_OK(Commit(realtime_commits, prepare_identifier)); + ASSERT_OK(writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(WithExpectedOffsets(expected_rows), actual_rows); + } + + void RunConcurrentPrepareTest(int32_t prepare_thread_count) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + + constexpr int64_t kBatchCount = 50; + constexpr int64_t kRowsPerBatch = 4; + std::vector expected_rows = MakeRows( + /*first_id=*/0, kBatchCount * kRowsPerBatch, /*partition=*/"p0"); + + std::atomic start{false}; + std::atomic writer_done{false}; + std::atomic ready_threads{0}; + std::atomic next_identifier{0}; + std::mutex result_mutex; + std::vector realtime_commits; + std::vector errors; + std::vector prepare_call_counts(prepare_thread_count, 0); + + auto record_error = [&](const Status& status) { + std::lock_guard lock(result_mutex); + errors.push_back(status.ToString()); + }; + auto append_commits = [&](std::vector&& commits) { + std::lock_guard lock(result_mutex); + realtime_commits.insert(realtime_commits.end(), + std::make_move_iterator(commits.begin()), + std::make_move_iterator(commits.end())); + }; + + std::thread write_thread([&]() { + ++ready_threads; + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + std::vector rows = MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, "p0"); + Result> batch_result = + MakeBatch(rows, /*partitioned=*/false); + if (!batch_result.ok()) { + record_error(batch_result.status()); + break; + } + Status status = writer->Write(std::move(batch_result).value()); + if (!status.ok()) { + record_error(status); + break; + } + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + writer_done.store(true, std::memory_order_release); + }); + + std::vector prepare_threads; + prepare_threads.reserve(prepare_thread_count); + for (int32_t thread_index = 0; thread_index < prepare_thread_count; ++thread_index) { + prepare_threads.emplace_back([&, thread_index]() { + ++ready_threads; + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + do { + int64_t identifier = next_identifier.fetch_add(1); + Result> result = + writer->PrepareCommitWithProgress(identifier); + ++prepare_call_counts[thread_index]; + if (!result.ok()) { + record_error(result.status()); + break; + } + append_commits(std::move(result).value()); + std::this_thread::sleep_for(std::chrono::microseconds(100)); + } while (!writer_done.load(std::memory_order_acquire)); + }); + } + + while (ready_threads.load(std::memory_order_acquire) < prepare_thread_count + 1) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + + write_thread.join(); + for (std::thread& thread : prepare_threads) { + thread.join(); + } + + ASSERT_TRUE(errors.empty()) << (errors.empty() ? "" : errors.front()); + for (int32_t call_count : prepare_call_counts) { + ASSERT_GT(call_count, 0); + } + FinalizeCommitAndCheck(writer.get(), std::move(realtime_commits), + next_identifier.fetch_add(1), std::move(expected_rows)); + } + + std::unique_ptr dir_; + std::string table_path_; + std::string commit_user_ = "realtime_commit_user"; + arrow::FieldVector fields_; + std::shared_ptr schema_; + std::map options_; +}; + +TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); +} + +TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { + CreateTable(/*partition_keys=*/{"pt"}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector expected_rows; + for (int64_t partition_index = 0; partition_index < 3; ++partition_index) { + std::string partition = "p" + std::to_string(partition_index); + std::vector rows = MakeRows(partition_index * 10, /*count=*/10, partition); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(batch))); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + } + FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, + std::move(expected_rows)); + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap committed_offsets, + ReadCommittedOffsets()); + ASSERT_EQ(3, committed_offsets.size()); + for (int64_t partition_index = 0; partition_index < 3; ++partition_index) { + std::string partition = "pt=p" + std::to_string(partition_index); + PartitionBucket partition_bucket(partition, /*bucket=*/0); + ASSERT_EQ(9, committed_offsets.at(partition_bucket)); + } +} + +TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { + CreateTable(/*partition_keys=*/{}); + + std::vector first_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, CreateRealtimeWriter()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_commits, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, first_commits.size()); + ASSERT_EQ(Range(0, 2), first_commits[0].offset_range); + ASSERT_OK(Commit(first_commits, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->Close()); + + std::vector second_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, CreateRealtimeWriter()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_commits.size()); + ASSERT_EQ(Range(3, 4), second_commits[0].offset_range); + + std::vector expected_rows = MakeRows(/*first_id=*/0, /*count=*/5, /*partition=*/"p0"); + FinalizeCommitAndCheck(second_writer.get(), std::move(second_commits), + /*prepare_identifier=*/1, std::move(expected_rows)); + + PartitionBucket partition_bucket("", /*bucket=*/0); + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap second_committed_offsets, + ReadCommittedOffsets()); + ASSERT_EQ(4, second_committed_offsets.at(partition_bucket)); +} + +TEST_F(RealtimeWriteInteTest, TestConcurrentWriteAndPrepareCommit) { + RunConcurrentPrepareTest(/*prepare_thread_count=*/1); +} + +TEST_F(RealtimeWriteInteTest, TestConcurrentWriteAndMultiplePrepareCommitThreads) { + RunConcurrentPrepareTest(/*prepare_thread_count=*/4); +} + +} // namespace paimon::test From 4007207c3dc566f353429325165ee1522a924179 Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sat, 1 Aug 2026 17:56:53 +0800 Subject: [PATCH 2/4] adjust offset file --- .../realtime/realtime_commit_progress.h | 5 +- include/paimon/realtime/realtime_context.h | 3 +- .../operation/abstract_file_store_write.cpp | 9 +- .../append_only_file_store_write.cpp | 15 ++-- .../append_only_file_store_write_test.cpp | 4 +- .../commit/realtime_snapshot_properties.cpp | 21 ++--- .../realtime_snapshot_properties_test.cpp | 84 +++++++++++-------- .../core/operation/file_store_commit_impl.cpp | 5 +- src/paimon/core/realtime/partition_bucket.h | 16 ++-- .../realtime/realtime_append_only_writer.cpp | 3 +- .../realtime/realtime_append_only_writer.h | 3 +- src/paimon/core/realtime/realtime_context.cpp | 7 +- test/inte/realtime_write_inte_test.cpp | 5 +- 13 files changed, 96 insertions(+), 84 deletions(-) diff --git a/include/paimon/realtime/realtime_commit_progress.h b/include/paimon/realtime/realtime_commit_progress.h index 6f346ca8..01a7ba48 100644 --- a/include/paimon/realtime/realtime_commit_progress.h +++ b/include/paimon/realtime/realtime_commit_progress.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include @@ -34,8 +35,8 @@ namespace paimon { struct PAIMON_EXPORT RealtimeCommitProgress { /// Paimon commit message generated from one sealed segment. std::shared_ptr commit_message; - /// Partition path such as `dt=2`, or an empty string for an unpartitioned table. - std::string partition; + /// Logical partition values, or an empty map for an unpartitioned table. + std::map partition; /// Bucket containing the sealed segment. int32_t bucket; /// Inclusive offset range represented by the commit message. diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index 8f4db839..fa07c0a2 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -50,7 +50,8 @@ class PAIMON_EXPORT RealtimeContext { /// Returns the stable indexer associated with a partition and bucket, creating it if needed. Result> GetOrCreateMemIndexer( - const std::string& partition, int32_t bucket, std::unique_ptr<::ArrowSchema> write_schema, + const std::map& partition, int32_t bucket, + std::unique_ptr<::ArrowSchema> write_schema, const std::map& options, const std::shared_ptr& memory_pool); diff --git a/src/paimon/core/operation/abstract_file_store_write.cpp b/src/paimon/core/operation/abstract_file_store_write.cpp index 14ff04be..fa8af413 100644 --- a/src/paimon/core/operation/abstract_file_store_write.cpp +++ b/src/paimon/core/operation/abstract_file_store_write.cpp @@ -31,7 +31,6 @@ #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/partition_bucket.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" #include "paimon/core/table/bucket_mode.h" @@ -315,9 +314,11 @@ Result> AbstractFileStoreWrite::PrepareRealt if (committable->IsEmpty()) { return Status::Invalid("sealed real-time segment produced an empty commit message"); } - PAIMON_ASSIGN_OR_RAISE(std::string partition, - file_store_path_factory_->GetPartitionString(snapshot.partition)); - partition = PartitionBucket::NormalizePartition(std::move(partition)); + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, file_store_path_factory_->GeneratePartitionVector( + snapshot.partition)); + std::map partition(partition_values.begin(), + partition_values.end()); result.push_back(RealtimeCommitProgress{committable, std::move(partition), snapshot.bucket, increment.GetRealtimeOffsetRange().value()}); } 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 2f4073af..8d43ae26 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include "arrow/c/bridge.h" @@ -237,13 +238,15 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( return std::shared_ptr(std::move(writer)); } - PAIMON_ASSIGN_OR_RAISE(std::string partition_string, - file_store_path_factory_->GetPartitionString(partition)); - partition_string = PartitionBucket::NormalizePartition(std::move(partition_string)); + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + 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())); int64_t next_offset = 0; - PartitionBucket partition_bucket(partition_string, bucket); + PartitionBucket partition_bucket(partition_map, bucket); auto offset_iter = realtime_committed_offsets_.find(partition_bucket); if (offset_iter != realtime_committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { @@ -252,8 +255,8 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( next_offset = offset_iter->second + 1; } return RealtimeAppendOnlyWriter::Create( - partition_string, bucket, std::move(c_write_schema), realtime_context_, writer, - write_schema_, realtime_write_schema_, options_.ToMap(), next_offset, pool_); + partition_map, bucket, std::move(c_write_schema), realtime_context_, writer, write_schema_, + realtime_write_schema_, options_.ToMap(), next_offset, pool_); } Result AppendOnlyFileStoreWrite::GetDataFileWriterFactory( 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 07b2d9cb..56bc0de7 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 @@ -310,7 +310,7 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestRealtimeWriteFlushesOffsetColumn) { ASSERT_OK_AND_ASSIGN(auto first_prepared, file_store_write->PrepareCommitWithProgress(/*commit_identifier=*/0)); ASSERT_EQ(1, first_prepared.size()); - ASSERT_EQ("", first_prepared[0].partition); + ASSERT_TRUE(first_prepared[0].partition.empty()); ASSERT_EQ(0, first_prepared[0].bucket); ASSERT_EQ(Range(0, 1), first_prepared[0].offset_range); std::shared_ptr first_file = OnlyNewFile({first_prepared[0].commit_message}); @@ -337,7 +337,7 @@ TEST_F(AppendOnlyFileStoreWriteTest, TestRealtimeWriteFlushesOffsetColumn) { ASSERT_OK_AND_ASSIGN(auto second_prepared, file_store_write->PrepareCommitWithProgress(/*commit_identifier=*/1)); ASSERT_EQ(1, second_prepared.size()); - ASSERT_EQ("", second_prepared[0].partition); + ASSERT_TRUE(second_prepared[0].partition.empty()); ASSERT_EQ(0, second_prepared[0].bucket); ASSERT_EQ(Range(2, 2), second_prepared[0].offset_range); std::shared_ptr second_file = OnlyNewFile({second_prepared[0].commit_message}); diff --git a/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp b/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp index b4b6f30f..6d2360ee 100644 --- a/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp +++ b/src/paimon/core/operation/commit/realtime_snapshot_properties.cpp @@ -38,7 +38,7 @@ class OffsetEntryJson { public: OffsetEntryJson() = default; - OffsetEntryJson(std::string partition, int32_t bucket, int64_t offset) + OffsetEntryJson(std::map partition, int32_t bucket, int64_t offset) : partition_(std::move(partition)), bucket_(bucket), offset_(offset) {} rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const { @@ -51,12 +51,13 @@ class OffsetEntryJson { } void FromJson(const rapidjson::Value& value) { - partition_ = RapidJsonUtil::DeserializeKeyValue(value, "partition"); + partition_ = RapidJsonUtil::DeserializeKeyValue>( + value, "partition"); bucket_ = RapidJsonUtil::DeserializeKeyValue(value, "bucket"); offset_ = RapidJsonUtil::DeserializeKeyValue(value, "offset"); } - const std::string& Partition() const { + const std::map& Partition() const { return partition_; } @@ -69,7 +70,7 @@ class OffsetEntryJson { } private: - std::string partition_; + std::map partition_; int32_t bucket_ = -1; int64_t offset_ = -1; }; @@ -95,10 +96,8 @@ class OffsetsJson { fmt::format("invalid bucket {} in offsets", partition_bucket.bucket)); } if (offset < 0) { - throw std::invalid_argument( - fmt::format("invalid offset {} for partition '{}' " - "bucket {}", - offset, partition_bucket.partition, partition_bucket.bucket)); + throw std::invalid_argument(fmt::format("invalid offset {} for bucket {}", offset, + partition_bucket.bucket)); } entries.emplace_back(partition_bucket.partition, partition_bucket.bucket, offset); } @@ -126,8 +125,7 @@ class OffsetsJson { PartitionBucket partition_bucket(entry.Partition(), entry.Bucket()); if (!offsets_.emplace(std::move(partition_bucket), entry.Offset()).second) { throw std::invalid_argument( - fmt::format("duplicate partition '{}' bucket {} in offsets", entry.Partition(), - entry.Bucket())); + fmt::format("duplicate partition-bucket {} in offsets", entry.Bucket())); } } } @@ -182,8 +180,7 @@ RealtimeSnapshotProperties::ValidateProgress(const std::vector::max() || commit.offset_range.from != previous_offset + 1) { return Status::Invalid(fmt::format( - "real-time commit offsets for partition '{}' bucket {} are not contiguous", - commit.partition, commit.bucket)); + "real-time commit offsets for bucket {} are not contiguous", commit.bucket)); } last_offsets[partition_bucket] = commit.offset_range.to; diff --git a/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp b/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp index 512e64ca..1180a568 100644 --- a/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp +++ b/src/paimon/core/operation/commit/realtime_snapshot_properties_test.cpp @@ -39,17 +39,22 @@ const std::string kTargetJson = R"({ "version": 1, "offsets": [ { - "partition": "", + "partition": {}, "bucket": 0, "offset": 7 }, { - "partition": "dt=2", + "partition": { + "dt": "2" + }, "bucket": 2, "offset": 9 }, { - "partition": "dt=2", + "partition": { + "dt": "a/b", + "region": "cn" + }, "bucket": 10, "offset": 12 } @@ -109,9 +114,9 @@ class RealtimeSnapshotPropertiesTest : public testing::Test { }; RealtimeSnapshotProperties::OffsetMap TargetOffsets() { - return {{PartitionBucket("", /*bucket=*/0), 7}, - {PartitionBucket("dt=2", /*bucket=*/2), 9}, - {PartitionBucket("dt=2", /*bucket=*/10), 12}}; + return {{PartitionBucket(/*partition=*/{}, /*bucket=*/0), 7}, + {PartitionBucket({{"dt", "2"}}, /*bucket=*/2), 9}, + {PartitionBucket({{"dt", "a/b"}, {"region", "cn"}}, /*bucket=*/10), 12}}; } } // namespace @@ -140,20 +145,21 @@ TEST_F(RealtimeSnapshotPropertiesTest, SerializeEmptyOffsets) { } TEST_F(RealtimeSnapshotPropertiesTest, SerializeRejectsInvalidOffsets) { - RealtimeSnapshotProperties::OffsetMap negative_bucket = {{PartitionBucket("dt=2", -1), 0}}; + RealtimeSnapshotProperties::OffsetMap negative_bucket = { + {PartitionBucket({{"dt", "2"}}, -1), 0}}; ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::SerializeOffsets(negative_bucket), "invalid bucket -1"); RealtimeSnapshotProperties::OffsetMap negative_offset = { - {PartitionBucket("dt=2", /*bucket=*/0), -1}}; + {PartitionBucket({{"dt", "2"}}, /*bucket=*/0), -1}}; ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::SerializeOffsets(negative_offset), "invalid offset -1"); } -TEST_F(RealtimeSnapshotPropertiesTest, NormalizePartitionBucketAndOffsetsDirectory) { - ASSERT_EQ("dt=2", PartitionBucket::NormalizePartition("dt=2/")); - PartitionBucket expected_partition_bucket("dt=2", 3); - ASSERT_EQ(expected_partition_bucket, PartitionBucket("dt=2/", /*bucket=*/3)); +TEST_F(RealtimeSnapshotPropertiesTest, PartitionBucketAndOffsetsDirectory) { + PartitionBucket expected_partition_bucket({{"dt", "a/b"}, {"region", "cn"}}, 3); + ASSERT_EQ(expected_partition_bucket, + PartitionBucket({{"dt", "a/b"}, {"region", "cn"}}, /*bucket=*/3)); ASSERT_EQ("/table/metadata", RealtimeSnapshotProperties::OffsetsDirectory("/table", "main")); ASSERT_EQ("/table/branch/branch-dev/metadata", RealtimeSnapshotProperties::OffsetsDirectory("/table", "dev")); @@ -202,18 +208,22 @@ TEST_F(RealtimeSnapshotPropertiesTest, DeserializeRejectsInvalidJson) { {R"({"version":2,"offsets":[]})", "unsupported offsets version 2"}, {R"({"version":1,"offsets":{}})", "value must be an array"}, {R"({"version":1,"offsets":[{"bucket":0,"offset":1}]})", "key must exist"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","offset":1}]})", "key must exist"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0}]})", "key must exist"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":"0","offset":1}]})", + {R"({"version":1,"offsets":[{"partition":{},"offset":1}]})", "key must exist"}, + {R"({"version":1,"offsets":[{"partition":{},"bucket":0}]})", "key must exist"}, + {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":1}]})", + "value must be an object"}, + {R"({"version":1,"offsets":[{"partition":{"dt":2},"bucket":0,"offset":1}]})", + "value must be string"}, + {R"({"version":1,"offsets":[{"partition":{},"bucket":"0","offset":1}]})", "value must be int"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":"1"}]})", + {R"({"version":1,"offsets":[{"partition":{},"bucket":0,"offset":"1"}]})", "value must be int64"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":-1,"offset":1}]})", + {R"({"version":1,"offsets":[{"partition":{},"bucket":-1,"offset":1}]})", "invalid bucket -1"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":-1}]})", + {R"({"version":1,"offsets":[{"partition":{},"bucket":0,"offset":-1}]})", "invalid offset -1"}, - {R"({"version":1,"offsets":[{"partition":"dt=2","bucket":0,"offset":1},{"partition":"dt=2/","bucket":0,"offset":2}]})", - "duplicate partition 'dt=2/' bucket 0"}}; + {R"({"version":1,"offsets":[{"partition":{"dt":"a/b"},"bucket":0,"offset":1},{"partition":{"dt":"a/b"},"bucket":0,"offset":2}]})", + "duplicate partition-bucket 0"}}; for (const auto& [json, expected_error] : invalid_json_cases) { SCOPED_TRACE(json); @@ -222,13 +232,13 @@ TEST_F(RealtimeSnapshotPropertiesTest, DeserializeRejectsInvalidJson) { } TEST_F(RealtimeSnapshotPropertiesTest, ValidateAndOrderProgress) { - PartitionBucket bucket0("dt=2", /*bucket=*/0); - PartitionBucket bucket1("dt=2", /*bucket=*/1); + PartitionBucket bucket0({{"dt", "2"}}, /*bucket=*/0); + PartitionBucket bucket1({{"dt", "2"}}, /*bucket=*/1); RealtimeSnapshotProperties::OffsetMap committed_offsets = {{bucket1, 4}}; std::vector commits = { - {/*commit_message=*/nullptr, "dt=2", 0, Range(2, 3)}, - {/*commit_message=*/nullptr, "dt=2", 1, Range(5, 6)}, - {/*commit_message=*/nullptr, "dt=2", 0, Range(0, 1)}}; + {/*commit_message=*/nullptr, {{"dt", "2"}}, 0, Range(2, 3)}, + {/*commit_message=*/nullptr, {{"dt", "2"}}, 1, Range(5, 6)}, + {/*commit_message=*/nullptr, {{"dt", "2"}}, 0, Range(0, 1)}}; ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::ValidatedCommitProgress validated_progress, RealtimeSnapshotProperties::ValidateProgress(commits, committed_offsets)); @@ -241,29 +251,31 @@ TEST_F(RealtimeSnapshotPropertiesTest, ValidateAndOrderProgress) { } TEST_F(RealtimeSnapshotPropertiesTest, ValidateProgressRejectsInvalidProgress) { - PartitionBucket bucket0("dt=2", /*bucket=*/0); + PartitionBucket bucket0({{"dt", "2"}}, /*bucket=*/0); RealtimeSnapshotProperties::OffsetMap committed_offsets = {{bucket0, 1}}; std::vector invalid_bucket = { - {/*commit_message=*/nullptr, "dt=2", -1, Range(0, 0)}}; + {/*commit_message=*/nullptr, {{"dt", "2"}}, -1, Range(0, 0)}}; ASSERT_NOK_WITH_MSG( RealtimeSnapshotProperties::ValidateProgress(invalid_bucket, /*committed_offsets=*/{}), "bucket -1 is invalid"); std::vector gap = { - {/*commit_message=*/nullptr, "dt=2", 0, Range(3, 4)}}; + {/*commit_message=*/nullptr, {{"dt", "2"}}, 0, Range(3, 4)}}; ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(gap, committed_offsets), "are not contiguous"); std::vector overlap = { - {/*commit_message=*/nullptr, "dt=2", 0, Range(1, 2)}}; + {/*commit_message=*/nullptr, {{"dt", "2"}}, 0, Range(1, 2)}}; ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(overlap, committed_offsets), "are not contiguous"); RealtimeSnapshotProperties::OffsetMap exhausted_offsets = { {bucket0, std::numeric_limits::max()}}; std::vector after_max = { - {/*commit_message=*/nullptr, "dt=2", 0, + {/*commit_message=*/nullptr, + {{"dt", "2"}}, + 0, Range(std::numeric_limits::max(), std::numeric_limits::max())}}; ASSERT_NOK_WITH_MSG(RealtimeSnapshotProperties::ValidateProgress(after_max, exhausted_offsets), "are not contiguous"); @@ -303,9 +315,9 @@ TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsWritesMergedProgress) { {RealtimeSnapshotProperties::kOffsetsKey, latest_offsets_path}}; RealtimeSnapshotProperties::OffsetMap delta_offsets = { - {PartitionBucket("", /*bucket=*/0), 5}, - {PartitionBucket("dt=2", /*bucket=*/2), 11}, - {PartitionBucket("dt=3", /*bucket=*/0), 4}}; + {PartitionBucket(/*partition=*/{}, /*bucket=*/0), 5}, + {PartitionBucket({{"dt", "2"}}, /*bucket=*/2), 11}, + {PartitionBucket({{"dt", "3"}}, /*bucket=*/0), 4}}; ASSERT_OK_AND_ASSIGN(std::string delta_json, RealtimeSnapshotProperties::SerializeOffsets(delta_offsets)); std::map properties = { @@ -323,8 +335,8 @@ TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsWritesMergedProgress) { RealtimeSnapshotProperties::ReadOffsets( std::optional(MakeSnapshot(merged)), file_system_)); RealtimeSnapshotProperties::OffsetMap expected = TargetOffsets(); - expected[PartitionBucket("dt=2", /*bucket=*/2)] = 11; - expected[PartitionBucket("dt=3", /*bucket=*/0)] = 4; + expected[PartitionBucket({{"dt", "2"}}, /*bucket=*/2)] = 11; + expected[PartitionBucket({{"dt", "3"}}, /*bucket=*/0)] = 4; ASSERT_EQ(expected, actual); } @@ -344,7 +356,7 @@ TEST_F(RealtimeSnapshotPropertiesTest, MergeEmptyDelta) { TEST_F(RealtimeSnapshotPropertiesTest, MergeOffsetsRequiresFileSystem) { RealtimeSnapshotProperties::OffsetMap delta_offsets = { - {PartitionBucket("dt=2", /*bucket=*/0), 1}}; + {PartitionBucket({{"dt", "2"}}, /*bucket=*/0), 1}}; ASSERT_OK_AND_ASSIGN(std::string delta_json, RealtimeSnapshotProperties::SerializeOffsets(delta_offsets)); std::map properties = { diff --git a/src/paimon/core/operation/file_store_commit_impl.cpp b/src/paimon/core/operation/file_store_commit_impl.cpp index 3c58979e..565990e6 100644 --- a/src/paimon/core/operation/file_store_commit_impl.cpp +++ b/src/paimon/core/operation/file_store_commit_impl.cpp @@ -901,9 +901,8 @@ Status FileStoreCommitImpl::CommitWithProgress( if (!commit_message) { return Status::Invalid("fail to cast real-time commit message to impl"); } - PAIMON_ASSIGN_OR_RAISE(std::string message_partition, - path_factory_->GetPartitionString(commit_message->Partition())); - message_partition = PartitionBucket::NormalizePartition(std::move(message_partition)); + std::map message_partition; + PAIMON_ASSIGN_OR_RAISE(message_partition, PartitionToMap(commit_message->Partition())); if (message_partition != realtime_commit.partition || commit_message->Bucket() != realtime_commit.bucket) { return Status::Invalid( diff --git a/src/paimon/core/realtime/partition_bucket.h b/src/paimon/core/realtime/partition_bucket.h index be04f9b6..e6c78b1a 100644 --- a/src/paimon/core/realtime/partition_bucket.h +++ b/src/paimon/core/realtime/partition_bucket.h @@ -17,22 +17,16 @@ #pragma once #include +#include #include #include -#include "paimon/common/utils/path_util.h" - namespace paimon { -/// Identifies one partition-bucket after normalizing the partition path. +/// Identifies one partition-bucket by its logical partition values. struct PartitionBucket { - PartitionBucket(std::string partition, int32_t bucket) - : partition(NormalizePartition(std::move(partition))), bucket(bucket) {} - - static std::string NormalizePartition(std::string partition) { - PathUtil::TrimLastDelim(&partition); - return partition; - } + PartitionBucket(std::map partition, int32_t bucket) + : partition(std::move(partition)), bucket(bucket) {} bool operator<(const PartitionBucket& other) const { if (partition != other.partition) { @@ -45,7 +39,7 @@ struct PartitionBucket { return partition == other.partition && bucket == other.bucket; } - std::string partition; + std::map partition; int32_t bucket = -1; }; diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 4226307c..cbc043a1 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -38,7 +38,8 @@ namespace paimon { Result> RealtimeAppendOnlyWriter::Create( - const std::string& partition, int32_t bucket, std::unique_ptr<::ArrowSchema> write_schema, + 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, diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h index 2560a049..c34d0cab 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.h +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -41,7 +41,8 @@ class RealtimeContext; class RealtimeAppendOnlyWriter : public BatchWriter { public: static Result> Create( - const std::string& partition, int32_t bucket, std::unique_ptr<::ArrowSchema> write_schema, + 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, diff --git a/src/paimon/core/realtime/realtime_context.cpp b/src/paimon/core/realtime/realtime_context.cpp index 9a9a0dcd..d1e792ec 100644 --- a/src/paimon/core/realtime/realtime_context.cpp +++ b/src/paimon/core/realtime/realtime_context.cpp @@ -34,7 +34,8 @@ class RealtimeContext::Impl { explicit Impl(const std::shared_ptr& factory) : factory_(factory) {} Result> GetOrCreateMemIndexer( - const std::string& partition, int32_t bucket, std::unique_ptr write_schema, + const std::map& partition, int32_t bucket, + std::unique_ptr write_schema, const std::map& options, const std::shared_ptr& memory_pool) { std::lock_guard lock(mutex_); @@ -79,8 +80,8 @@ RealtimeContext::RealtimeContext(std::unique_ptr&& impl) : impl_(std::move RealtimeContext::~RealtimeContext() = default; Result> RealtimeContext::GetOrCreateMemIndexer( - const std::string& partition, int32_t bucket, std::unique_ptr write_schema, - const std::map& options, + const std::map& partition, int32_t bucket, + std::unique_ptr write_schema, const std::map& options, const std::shared_ptr& memory_pool) { return impl_->GetOrCreateMemIndexer(partition, bucket, std::move(write_schema), options, memory_pool); diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 8e5736e7..2947de9f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -371,7 +371,8 @@ TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { ReadCommittedOffsets()); ASSERT_EQ(3, committed_offsets.size()); for (int64_t partition_index = 0; partition_index < 3; ++partition_index) { - std::string partition = "pt=p" + std::to_string(partition_index); + std::map partition = { + {"pt", "p" + std::to_string(partition_index)}}; PartitionBucket partition_bucket(partition, /*bucket=*/0); ASSERT_EQ(9, committed_offsets.at(partition_bucket)); } @@ -406,7 +407,7 @@ TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { FinalizeCommitAndCheck(second_writer.get(), std::move(second_commits), /*prepare_identifier=*/1, std::move(expected_rows)); - PartitionBucket partition_bucket("", /*bucket=*/0); + PartitionBucket partition_bucket(/*partition=*/{}, /*bucket=*/0); ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap second_committed_offsets, ReadCommittedOffsets()); ASSERT_EQ(4, second_committed_offsets.at(partition_bucket)); From dfa49f50d442e4638aec959fb4952ac4aa7235fe Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sun, 2 Aug 2026 12:29:12 +0800 Subject: [PATCH 3/4] support union read and reclaim --- include/paimon/file_store_write.h | 6 + include/paimon/realtime/mem_indexer.h | 51 +- include/paimon/realtime/realtime_context.h | 44 +- include/paimon/scan_context.h | 11 + src/paimon/CMakeLists.txt | 1 + .../append_only_file_store_write.cpp | 36 +- .../operation/append_only_file_store_write.h | 4 + .../core/operation/file_store_commit_impl.cpp | 29 +- .../core/operation/file_store_commit_impl.h | 8 +- .../operation/file_store_commit_impl_test.cpp | 6 +- .../core/operation/file_store_write.cpp | 7 +- src/paimon/core/operation/scan_context.cpp | 15 +- .../core/realtime/arrow_mem_indexer.cpp | 225 +++- src/paimon/core/realtime/arrow_mem_indexer.h | 18 + .../realtime/arrow_mem_indexer_factory.cpp | 2 +- src/paimon/core/realtime/realtime_context.cpp | 53 + .../table/source/append_only_table_read.cpp | 48 + .../table/source/append_only_table_read.h | 2 + src/paimon/core/table/source/realtime_split.h | 83 ++ .../core/table/source/realtime_table_scan.cpp | 172 +++ .../core/table/source/realtime_table_scan.h | 80 ++ src/paimon/core/table/source/table_scan.cpp | 40 + test/inte/realtime_write_inte_test.cpp | 987 ++++++++++++++++-- 23 files changed, 1779 insertions(+), 149 deletions(-) create mode 100644 src/paimon/core/table/source/realtime_split.h create mode 100644 src/paimon/core/table/source/realtime_table_scan.cpp create mode 100644 src/paimon/core/table/source/realtime_table_scan.h diff --git a/include/paimon/file_store_write.h b/include/paimon/file_store_write.h index d79b2595..fdc172c3 100644 --- a/include/paimon/file_store_write.h +++ b/include/paimon/file_store_write.h @@ -103,6 +103,12 @@ class PAIMON_EXPORT FileStoreWrite { virtual Result> PrepareCommitWithProgress( int64_t commit_identifier); + /// Refreshes a real-time writer after `snapshot_id` has committed successfully. + /// + /// The writer loads the snapshot's partition-bucket offsets and releases sealed memory that is + /// fully covered by disk. Calling this method on a non-real-time writer returns an error. + virtual Status RefreshCommittedSnapshot(int64_t snapshot_id); + virtual std::shared_ptr GetMetrics() const = 0; virtual Status Close() = 0; }; diff --git a/include/paimon/realtime/mem_indexer.h b/include/paimon/realtime/mem_indexer.h index 18cd13ea..3e318602 100644 --- a/include/paimon/realtime/mem_indexer.h +++ b/include/paimon/realtime/mem_indexer.h @@ -34,6 +34,7 @@ struct ArrowSchema; namespace paimon { class MemoryPool; +class Predicate; /// A record batch and the contiguous offset range assigned to its rows. /// @@ -59,6 +60,34 @@ class PAIMON_EXPORT RealtimeSegmentHandle { virtual Range GetOffsetRange() const = 0; }; +/// Opaque immutable view of the rows visible from one `MemIndexer`. +/// +/// A view pins all referenced resources until the readers created from it are closed. Later +/// writes, seals, and committed-offset reclamation do not change the contents of an existing view. +class PAIMON_EXPORT MemReadView { + public: + virtual ~MemReadView() = default; + + /// Returns the inclusive offset range visible in this view, or no range when it is empty. + virtual std::optional GetOffsetRange() const = 0; +}; + +/// Parameters used by a `MemIndexer` to create readers for a query. +struct PAIMON_EXPORT MemQueryContext { + /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added. + /// + /// The schema uses the Arrow C Data Interface and is valid only during + /// `CreateQueryReaders`. An implementation may consume it with an Arrow importer. + ::ArrowSchema* read_schema; + /// Predicate using field indexes from `read_schema`. + std::shared_ptr predicate; + /// Whether exact predicate filtering is enabled for the returned readers. + /// + /// Keep this disabled for primary-key merge-on-read. Filtering memory before PK merge may + /// remove the newest row and incorrectly expose an older disk row. + bool enable_predicate_filter; +}; + /// Plugin interface for buffering real-time writes before Paimon data-file generation. /// /// Paimon serializes calls to `Write` and `SealForCommit` for the same indexer. After sealing, @@ -87,7 +116,27 @@ class PAIMON_EXPORT MemIndexer { virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; - /// Returns the number of bytes currently retained by the building segment. + /// Acquires an immutable view containing the current sealed and building rows. + /// + /// This method may be called concurrently with query-reader creation and reclamation. It must + /// also provide a consistent snapshot when a write or seal is in progress. + virtual Result> AcquireReadView() = 0; + + /// Creates readers over rows in `view` whose offsets are greater than + /// `offset_lower_exclusive`. + /// + /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by + /// `context.read_schema` except a duplicate `_VALUE_KIND`. `_OFFSET` is returned when requested + /// by the read schema. Concatenating all returned readers must produce every matching row once. + virtual Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_lower_exclusive, + const MemQueryContext& context) = 0; + + /// Releases this indexer's ownership of sealed segments fully covered by the committed offset. + /// Existing read views continue to keep their referenced resources alive. + virtual Status Reclaim(int64_t committed_offset) = 0; + + /// Returns the number of bytes currently retained by building and sealed segments. virtual uint64_t GetMemoryUsage() const = 0; /// Releases resources owned by this indexer and rejects subsequent writes or seals. diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index fa07c0a2..f530e0b5 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -20,6 +20,7 @@ #include #include #include +#include #include "paimon/result.h" #include "paimon/visibility.h" @@ -30,13 +31,36 @@ namespace paimon { class MemIndexer; class MemIndexerFactory; +class MemReadView; class MemoryPool; +/// One partition-bucket and the immutable plugin view captured for a table scan. +struct PAIMON_EXPORT RealtimePartitionBucketView { + /// Logical partition values, before partition-path escaping. + std::map partition; + /// Fixed bucket id. + int32_t bucket; + /// Plugin instance that creates readers from `read_view`. + std::shared_ptr indexer; + /// Immutable rows pinned for one query plan. + std::shared_ptr read_view; +}; + +/// Committed progress for one partition-bucket loaded from a Paimon snapshot. +struct PAIMON_EXPORT RealtimePartitionBucketOffset { + /// Logical partition values, before partition-path escaping. + std::map partition; + /// Fixed bucket id. + int32_t bucket; + /// Largest `_OFFSET` committed for this partition-bucket. + int64_t offset; +}; + /// Shared context that owns the `MemIndexer` instances used by a real-time writer. /// -/// Applications attach one context to `WriteContext`. The context uses either the default Arrow -/// implementation or an application-provided factory and keeps each created indexer available -/// across multiple writes and prepare-commit operations. +/// Applications share one context between `WriteContext` and `ScanContext`. The context uses +/// either the default Arrow implementation or an application-provided factory and keeps each +/// created indexer available across writes, prepare-commit operations, and process-local reads. class PAIMON_EXPORT RealtimeContext { public: /// Creates a context backed by Paimon's default Arrow `MemIndexer`. @@ -55,6 +79,20 @@ class PAIMON_EXPORT RealtimeContext { const std::map& options, const std::shared_ptr& memory_pool); + /// Captures an immutable read view from every currently registered indexer. + /// + /// The indexer registry is fixed during this call and each returned plugin view is stable. New + /// partition-buckets registered after this call are not visible in that query. + Result> AcquireReadViews(); + + /// Advances the committed progress visible to the registered memory indexers. + /// + /// Calls are idempotent for the same snapshot and must advance snapshot ids monotonically. + /// Each indexer may release sealed data covered by its committed offset. Existing read views + /// continue to pin referenced resources until their readers are closed. + Status AdvanceCommittedProgress( + int64_t snapshot_id, const std::vector& committed_offsets); + ~RealtimeContext(); private: diff --git a/include/paimon/scan_context.h b/include/paimon/scan_context.h index 2dea6ac5..dc780ad7 100644 --- a/include/paimon/scan_context.h +++ b/include/paimon/scan_context.h @@ -48,6 +48,7 @@ class PAIMON_EXPORT ScanContext { ScanContext(const std::string& path, bool is_streaming_mode, std::optional limit, const std::shared_ptr& scan_filter, const std::shared_ptr& global_index_result, + const std::shared_ptr& realtime_context, const std::shared_ptr& memory_pool, const std::shared_ptr& executor, const std::shared_ptr& specific_file_system, @@ -87,6 +88,11 @@ class PAIMON_EXPORT ScanContext { return global_index_result_; } + /// Returns the optional process-local context used to plan real-time memory reads. + std::shared_ptr GetRealtimeContext() const { + return realtime_context_; + } + std::shared_ptr GetSpecificFileSystem() const { return specific_file_system_; } @@ -105,6 +111,7 @@ class PAIMON_EXPORT ScanContext { std::optional limit_; std::shared_ptr scan_filters_; std::shared_ptr global_index_result_; + std::shared_ptr realtime_context_; std::shared_ptr memory_pool_; std::shared_ptr executor_; std::shared_ptr specific_file_system_; @@ -163,6 +170,10 @@ class PAIMON_EXPORT ScanContextBuilder { ScanContextBuilder& SetGlobalIndexResult( const std::shared_ptr& global_index_result); + /// Enables process-local union reads with the memory indexers owned by `realtime_context`. + ScanContextBuilder& WithRealtimeContext( + const std::shared_ptr& realtime_context); + /// The options added or set in `ScanContextBuilder` have high priority and will be merged with /// the options in table schema. ScanContextBuilder& AddOption(const std::string& key, const std::string& value); diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index b236676f..8e705bfb 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -382,6 +382,7 @@ set(PAIMON_CORE_SRCS core/table/source/merge_tree_split_generator.cpp core/table/source/data_evolution_split_generator.cpp core/table/source/plan_impl.cpp + core/table/source/realtime_table_scan.cpp core/table/source/snapshot/snapshot_reader.cpp core/table/source/startup_mode.cpp core/table/source/table_read.cpp 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 8d43ae26..d76cc759 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -106,6 +106,29 @@ AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( AppendOnlyFileStoreWrite::~AppendOnlyFileStoreWrite() = default; +Status AppendOnlyFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + if (!realtime_context_) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeSnapshotProperties::OffsetMap committed_offsets, + RealtimeSnapshotProperties::ReadOffsets(std::optional(std::move(snapshot)), + options_.GetFileSystem())); + std::vector progress; + progress.reserve(committed_offsets.size()); + for (const auto& [partition_bucket, offset] : committed_offsets) { + progress.push_back(RealtimePartitionBucketOffset{partition_bucket.partition, + partition_bucket.bucket, offset}); + } + PAIMON_RETURN_NOT_OK(realtime_context_->AdvanceCommittedProgress(snapshot_id, progress)); + { + std::lock_guard lock(realtime_offsets_mutex_); + realtime_committed_offsets_ = std::move(committed_offsets); + } + return Status::OK(); +} + Result> AppendOnlyFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { PAIMON_ASSIGN_OR_RAISE( @@ -247,12 +270,15 @@ Result> AppendOnlyFileStoreWrite::CreateWriter( PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*write_schema_, c_write_schema.get())); int64_t next_offset = 0; PartitionBucket partition_bucket(partition_map, bucket); - auto offset_iter = realtime_committed_offsets_.find(partition_bucket); - if (offset_iter != realtime_committed_offsets_.end()) { - if (offset_iter->second == std::numeric_limits::max()) { - return Status::Invalid("real-time offset has reached INT64_MAX"); + { + std::lock_guard lock(realtime_offsets_mutex_); + auto offset_iter = realtime_committed_offsets_.find(partition_bucket); + if (offset_iter != realtime_committed_offsets_.end()) { + if (offset_iter->second == std::numeric_limits::max()) { + return Status::Invalid("real-time offset has reached INT64_MAX"); + } + next_offset = offset_iter->second + 1; } - next_offset = offset_iter->second + 1; } return RealtimeAppendOnlyWriter::Create( partition_map, bucket, std::move(c_write_schema), realtime_context_, writer, write_schema_, 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 442af906..788370e0 100644 --- a/src/paimon/core/operation/append_only_file_store_write.h +++ b/src/paimon/core/operation/append_only_file_store_write.h @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -87,6 +88,8 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr& executor, const std::shared_ptr& pool); ~AppendOnlyFileStoreWrite() override; + Status RefreshCommittedSnapshot(int64_t snapshot_id) override; + /// Rewrites the given files into new compacted files. /// /// @param partition The partition of the files. @@ -130,6 +133,7 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { std::optional> write_cols_; std::shared_ptr realtime_write_schema_; std::shared_ptr realtime_context_; + mutable std::mutex realtime_offsets_mutex_; RealtimeSnapshotProperties::OffsetMap realtime_committed_offsets_; std::unique_ptr logger_; }; diff --git a/src/paimon/core/operation/file_store_commit_impl.cpp b/src/paimon/core/operation/file_store_commit_impl.cpp index 565990e6..8ae24482 100644 --- a/src/paimon/core/operation/file_store_commit_impl.cpp +++ b/src/paimon/core/operation/file_store_commit_impl.cpp @@ -401,7 +401,8 @@ Result FileStoreCommitImpl::FilterAndCommit( if (!retry_committables.empty()) { PAIMON_RETURN_NOT_OK(CheckFilesExistence(retry_committables)); for (const auto& committable : retry_committables) { - PAIMON_RETURN_NOT_OK(Commit(committable, /*check_append_files=*/true)); + PAIMON_RETURN_NOT_OK(Commit(committable, /*check_append_files=*/true, + /*retry_on_conflict=*/true)); } } return retry_committables.size(); @@ -707,7 +708,8 @@ Status FileStoreCommitImpl::ExecuteOverwrite( TryCommit(changes->compact_table_files, /*changelog_files=*/{}, changes->compact_index_files, identifier, watermark, committable->Properties(), Snapshot::CommitKind::Compact(), - /*detect_conflicts=*/true)); + /*detect_conflicts=*/true, + /*retry_on_conflict=*/true)); *attempt += cnt; *generated_snapshot += 1; } @@ -808,11 +810,12 @@ Result FileStoreCommitImpl::TryOverwrite( std::shared_ptr changes_provider = commit_scanner_->OverwriteChangesProvider(partitions, changes, index_entries); return TryCommit(changes_provider, commit_identifier, watermark, properties, - Snapshot::CommitKind::Overwrite(), /*detect_conflicts=*/true); + Snapshot::CommitKind::Overwrite(), /*detect_conflicts=*/true, + /*retry_on_conflict=*/true); } Status FileStoreCommitImpl::Commit(const std::shared_ptr& committable, - bool check_append_files) { + bool check_append_files, bool retry_on_conflict) { PAIMON_LOG_INFO(logger_, "Ready to commit to table %s, number of commit messages: %zu", table_name_.c_str(), committable->FileCommittables().size()); PAIMON_ASSIGN_OR_RAISE(std::string committable_str, committable->ToString()); @@ -853,7 +856,7 @@ Status FileStoreCommitImpl::Commit(const std::shared_ptr& c int32_t cnt, TryCommit(changes.append_table_files, changes.append_changelog, changes.append_index_files, committable->Identifier(), committable->Watermark(), committable->Properties(), commit_kind, - check_append_files)); + check_append_files, retry_on_conflict)); attempt += cnt; generated_snapshot += 1; } @@ -864,7 +867,7 @@ Status FileStoreCommitImpl::Commit(const std::shared_ptr& c changes.compact_index_files, committable->Identifier(), committable->Watermark(), committable->Properties(), Snapshot::CommitKind::Compact(), - /*detect_conflicts=*/true)); + /*detect_conflicts=*/true, retry_on_conflict)); attempt += cnt; generated_snapshot += 1; } @@ -876,7 +879,7 @@ Status FileStoreCommitImpl::Commit( std::optional watermark) { std::shared_ptr committable = CreateManifestCommittable(identifier, commit_messages, watermark, /*properties=*/{}); - return Commit(committable, /*check_append_files=*/false); + return Commit(committable, /*check_append_files=*/false, /*retry_on_conflict=*/true); } Status FileStoreCommitImpl::CommitWithProgress( @@ -919,7 +922,7 @@ Status FileStoreCommitImpl::CommitWithProgress( } std::shared_ptr committable = CreateManifestCommittable(identifier, commit_messages, watermark, properties); - return Commit(committable, /*check_append_files=*/false); + return Commit(committable, /*check_append_files=*/false, /*retry_on_conflict=*/false); } Result FileStoreCommitImpl::TryCommit(const std::vector& delta_files, @@ -928,17 +931,17 @@ Result FileStoreCommitImpl::TryCommit(const std::vector& int64_t identifier, std::optional watermark, const std::map& properties, Snapshot::CommitKind commit_kind, - bool detect_conflicts) { + bool detect_conflicts, bool retry_on_conflict) { std::shared_ptr changes_provider = CommitChangesProvider::Provider(delta_files, changelog_files, index_entries); return TryCommit(changes_provider, identifier, watermark, properties, commit_kind, - detect_conflicts); + detect_conflicts, retry_on_conflict); } Result FileStoreCommitImpl::TryCommit( const std::shared_ptr& changes_provider, int64_t identifier, std::optional watermark, const std::map& properties, - Snapshot::CommitKind commit_kind, bool detect_conflicts) { + Snapshot::CommitKind commit_kind, bool detect_conflicts, bool retry_on_conflict) { int32_t retry_count = 0; int64_t start_millis = DateTimeUtils::GetCurrentUTCTimeUs() / 1000; while (true) { @@ -954,6 +957,10 @@ Result FileStoreCommitImpl::TryCommit( if (commit_success) { break; } + if (!retry_on_conflict) { + // TODO(xinyu.lxy): Support failure recovery and idempotent retry for real-time commits. + return Status::Invalid("real-time commit failed due to snapshot conflict"); + } int64_t current_millis = DateTimeUtils::GetCurrentUTCTimeUs() / 1000; if (current_millis - start_millis > options_.GetCommitTimeout() || retry_count >= options_.GetCommitMaxRetries()) { diff --git a/src/paimon/core/operation/file_store_commit_impl.h b/src/paimon/core/operation/file_store_commit_impl.h index 246455c8..f48d7281 100644 --- a/src/paimon/core/operation/file_store_commit_impl.h +++ b/src/paimon/core/operation/file_store_commit_impl.h @@ -147,7 +147,7 @@ class FileStoreCommitImpl : public FileStoreCommit { private: Status Commit(const std::shared_ptr& manifest_committable, - bool check_append_files); + bool check_append_files, bool retry_on_conflict); Result TryOverwrite(const std::vector>& partition, const std::vector& changes, @@ -188,12 +188,14 @@ class FileStoreCommitImpl : public FileStoreCommit { const std::vector& index_entries, int64_t identifier, std::optional watermark, const std::map& properties, - Snapshot::CommitKind commit_kind, bool detect_conflicts); + Snapshot::CommitKind commit_kind, bool detect_conflicts, + bool retry_on_conflict); Result TryCommit(const std::shared_ptr& changes_provider, int64_t identifier, std::optional watermark, const std::map& properties, - Snapshot::CommitKind commit_kind, bool detect_conflicts); + Snapshot::CommitKind commit_kind, bool detect_conflicts, + bool retry_on_conflict); Result TryCommitOnce(const std::vector& delta_files, const std::vector& changelog_files, diff --git a/src/paimon/core/operation/file_store_commit_impl_test.cpp b/src/paimon/core/operation/file_store_commit_impl_test.cpp index 2186dc88..2ff00672 100644 --- a/src/paimon/core/operation/file_store_commit_impl_test.cpp +++ b/src/paimon/core/operation/file_store_commit_impl_test.cpp @@ -1773,7 +1773,8 @@ TEST_F(FileStoreCommitImplTest, TestFilterCommitted) { ASSERT_EQ(filtered_committables.size(), committables.size()); // Test with a previous snapshot - ASSERT_OK(commit_impl->Commit(committable, /*check_append_files=*/false)); + ASSERT_OK(commit_impl->Commit(committable, /*check_append_files=*/false, + /*retry_on_conflict=*/true)); ASSERT_OK_AND_ASSIGN(filtered_committables, commit_impl->FilterCommitted(committables)); ASSERT_EQ(filtered_committables.size(), 0); } @@ -1812,7 +1813,8 @@ TEST_F(FileStoreCommitImplTest, TestFilterCommittedWithMultipleCommittables) { ASSERT_EQ(filtered_committables.size(), committables.size()); // Test with a previous snapshot - ASSERT_OK(commit_impl->Commit(committable1, /*check_append_files=*/false)); + ASSERT_OK(commit_impl->Commit(committable1, /*check_append_files=*/false, + /*retry_on_conflict=*/true)); ASSERT_OK_AND_ASSIGN(filtered_committables, commit_impl->FilterCommitted(committables)); ASSERT_EQ(filtered_committables.size(), 1); ASSERT_EQ(filtered_committables[0]->Identifier(), committable2->Identifier()); diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 6e79a3f2..a334aead 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -60,6 +60,10 @@ Result> FileStoreWrite::PrepareCommitWithPro return Status::Invalid("prepare commit with progress requires a real-time writer"); } +Status FileStoreWrite::RefreshCommittedSnapshot(int64_t) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); +} + Result> FileStoreWrite::Create(std::unique_ptr ctx) { if (ctx == nullptr) { return Status::Invalid("write context is null pointer"); @@ -177,12 +181,13 @@ Result> FileStoreWrite::Create(std::unique_ptr(index_file_handler); } - return std::make_unique( + auto file_store_write = std::make_unique( 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->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), ctx->GetRealtimeContext(), realtime_committed_offsets, ctx->GetExecutor(), ctx->GetMemoryPool()); + return std::unique_ptr(std::move(file_store_write)); } else { // pk table if (ctx->GetRealtimeContext()) { diff --git a/src/paimon/core/operation/scan_context.cpp b/src/paimon/core/operation/scan_context.cpp index b59fa098..16a80731 100644 --- a/src/paimon/core/operation/scan_context.cpp +++ b/src/paimon/core/operation/scan_context.cpp @@ -32,6 +32,7 @@ ScanContext::ScanContext(const std::string& path, bool is_streaming_mode, std::optional limit, const std::shared_ptr& scan_filter, const std::shared_ptr& global_index_result, + const std::shared_ptr& realtime_context, const std::shared_ptr& memory_pool, const std::shared_ptr& executor, const std::shared_ptr& specific_file_system, @@ -43,6 +44,7 @@ ScanContext::ScanContext(const std::string& path, bool is_streaming_mode, limit_(limit), scan_filters_(scan_filter), global_index_result_(global_index_result), + realtime_context_(realtime_context), memory_pool_(memory_pool), executor_(executor), specific_file_system_(specific_file_system), @@ -63,6 +65,7 @@ class ScanContextBuilder::Impl { partition_filters_.clear(); predicates_.reset(); global_index_result_.reset(); + realtime_context_.reset(); memory_pool_ = GetDefaultPool(); executor_ = CreateDefaultExecutor(); specific_file_system_.reset(); @@ -79,6 +82,7 @@ class ScanContextBuilder::Impl { std::vector> partition_filters_; std::shared_ptr predicates_; std::shared_ptr global_index_result_; + std::shared_ptr realtime_context_; std::shared_ptr memory_pool_ = GetDefaultPool(); std::shared_ptr executor_ = CreateDefaultExecutor(); std::shared_ptr specific_file_system_; @@ -124,6 +128,12 @@ ScanContextBuilder& ScanContextBuilder::SetGlobalIndexResult( return *this; } +ScanContextBuilder& ScanContextBuilder::WithRealtimeContext( + const std::shared_ptr& realtime_context) { + impl_->realtime_context_ = realtime_context; + return *this; +} + ScanContextBuilder& ScanContextBuilder::SetOptions( const std::map& options) { impl_->options_ = options; @@ -172,8 +182,9 @@ Result> ScanContextBuilder::Finish() { impl_->path_, impl_->is_streaming_mode_, impl_->limit_, std::make_shared(impl_->predicates_, impl_->partition_filters_, impl_->bucket_filter_), - impl_->global_index_result_, impl_->memory_pool_, impl_->executor_, - impl_->specific_file_system_, impl_->table_schema_, impl_->options_, impl_->cache_); + impl_->global_index_result_, impl_->realtime_context_, impl_->memory_pool_, + impl_->executor_, impl_->specific_file_system_, impl_->table_schema_, impl_->options_, + impl_->cache_); impl_->Reset(); return ctx; } diff --git a/src/paimon/core/realtime/arrow_mem_indexer.cpp b/src/paimon/core/realtime/arrow_mem_indexer.cpp index 136f0e2e..86de7d5e 100644 --- a/src/paimon/core/realtime/arrow_mem_indexer.cpp +++ b/src/paimon/core/realtime/arrow_mem_indexer.cpp @@ -16,14 +16,19 @@ #include "paimon/core/realtime/arrow_mem_indexer.h" +#include +#include #include #include "arrow/api.h" #include "arrow/c/bridge.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/reader/complete_row_kind_batch_reader.h" +#include "paimon/common/reader/predicate_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/core/utils/nested_projection_utils.h" #include "paimon/macros.h" namespace paimon { @@ -60,11 +65,41 @@ class ArrowMemIndexer::Segment : public RealtimeSegmentHandle { return batches_; } + uint64_t GetMemoryUsage() const { + uint64_t result = 0; + for (const StoredBatch& batch : batches_) { + result += batch.memory_usage; + } + return result; + } + private: Range offset_range_; std::vector batches_; }; +class ArrowMemIndexer::ReadView : public MemReadView { + public: + explicit ReadView(std::vector&& batches) : batches_(std::move(batches)) { + if (!batches_.empty()) { + offset_range_ = + Range(batches_.front().offset_range.from, batches_.back().offset_range.to); + } + } + + std::optional GetOffsetRange() const override { + return offset_range_; + } + + const std::vector& GetBatches() const { + return batches_; + } + + private: + std::vector batches_; + std::optional offset_range_; +}; + class ArrowMemIndexer::CommitBatchReader : public BatchReader { public: CommitBatchReader(const std::shared_ptr& segment, @@ -125,9 +160,110 @@ class ArrowMemIndexer::CommitBatchReader : public BatchReader { size_t next_batch_ = 0; }; +class ArrowMemIndexer::QueryBatchReader : public BatchReader { + public: + QueryBatchReader(const std::shared_ptr& view, int64_t offset_lower_exclusive, + const std::shared_ptr& read_schema, + const std::shared_ptr& arrow_pool) + : view_(view), + offset_lower_exclusive_(offset_lower_exclusive), + read_schema_(read_schema), + arrow_pool_(arrow_pool), + metrics_(std::make_shared()) {} + + Result NextBatch() override { + return Status::Invalid( + "paimon inner reader ArrowMemIndexer::QueryBatchReader 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. + if (offset_lower_exclusive_ == std::numeric_limits::max()) { + return MakeEofBatchWithBitmap(); + } + while (view_ && next_batch_ < view_->GetBatches().size()) { + const StoredBatch& stored = view_->GetBatches()[next_batch_++]; + if (stored.offset_range.to <= offset_lower_exclusive_) { + continue; + } + int64_t begin = + std::max(0, offset_lower_exclusive_ + 1 - stored.offset_range.from); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, BuildOutput(stored)); + RoaringBitmap32 candidate_rows; + candidate_rows.AddRange(static_cast(begin), + static_cast(stored.data->length())); + auto c_array = std::make_unique(); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*output, c_array.get(), c_schema.get())); + return ReadBatchWithBitmap(ReadBatch(std::move(c_array), std::move(c_schema)), + std::move(candidate_rows)); + } + return MakeEofBatchWithBitmap(); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + + void Close() override { + view_.reset(); + } + + private: + Result> BuildOutput(const StoredBatch& stored) { + std::shared_ptr source = stored.data; + if (read_schema_->GetFieldByName(SpecialFields::Offset().Name())) { + int64_t row_count = stored.data->length(); + arrow::Int64Builder offset_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offset_builder.Reserve(row_count)); + for (int64_t i = 0; i < row_count; ++i) { + offset_builder.UnsafeAppend(stored.offset_range.from + i); + } + std::shared_ptr offsets; + PAIMON_RETURN_NOT_OK_FROM_ARROW(offset_builder.Finish(&offsets)); + + arrow::ArrayVector source_arrays = {offsets}; + source_arrays.insert(source_arrays.end(), stored.data->fields().begin(), + stored.data->fields().end()); + arrow::FieldVector source_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::Offset())}; + source_fields.insert(source_fields.end(), stored.data->struct_type()->fields().begin(), + stored.data->struct_type()->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr source_with_offset, + arrow::StructArray::Make(source_arrays, source_fields, stored.data->null_bitmap(), + stored.data->null_count(), stored.data->offset())); + source = std::move(source_with_offset); + } + + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr projected, + NestedProjectionUtils::AlignArrayToReadType( + source, arrow::struct_(read_schema_->fields()), arrow_pool_.get())); + std::shared_ptr projected_struct = + std::dynamic_pointer_cast(projected); + if (!projected_struct) { + return Status::Invalid("memory query projection did not produce a StructArray"); + } + return projected_struct; + } + + private: + std::shared_ptr view_; + int64_t offset_lower_exclusive_; + std::shared_ptr read_schema_; + std::shared_ptr arrow_pool_; + std::shared_ptr metrics_; + size_t next_batch_ = 0; +}; + ArrowMemIndexer::ArrowMemIndexer(const std::shared_ptr& write_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool) - : write_schema_(write_schema), arrow_pool_(arrow_pool) {} + : write_schema_(write_schema), memory_pool_(memory_pool), arrow_pool_(arrow_pool) {} Status ArrowMemIndexer::Write(RealtimeWriteBatch&& write_batch) { if (!write_batch.batch) { @@ -137,6 +273,10 @@ Status ArrowMemIndexer::Write(RealtimeWriteBatch&& write_batch) { if (write_batch.offset_range.Count() != row_count) { return Status::Invalid("real-time offset range does not match batch row count"); } + 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, @@ -147,15 +287,18 @@ Status ArrowMemIndexer::Write(RealtimeWriteBatch&& write_batch) { return Status::Invalid("real-time write data is not a StructArray"); } + std::lock_guard lock(mutex_); if (closed_) { return Status::Invalid("mem indexer is closed"); } if (building_range_ && write_batch.offset_range.from != building_range_->to + 1) { return Status::Invalid("real-time offset ranges must be contiguous"); } - building_memory_usage_ += GetArrayMemoryUsage(struct_array->data()); - building_batches_.push_back( - StoredBatch{std::move(struct_array), write_batch.batch->GetRowKind()}); + uint64_t memory_usage = GetArrayMemoryUsage(struct_array->data()); + building_memory_usage_ += memory_usage; + building_batches_.push_back(StoredBatch{std::move(struct_array), + write_batch.batch->GetRowKind(), + write_batch.offset_range, memory_usage}); if (!building_range_) { building_range_ = write_batch.offset_range; } else { @@ -165,6 +308,7 @@ Status ArrowMemIndexer::Write(RealtimeWriteBatch&& write_batch) { } Result>> ArrowMemIndexer::SealForCommit() { + std::lock_guard lock(mutex_); if (closed_) { return Status::Invalid("mem indexer is closed"); } @@ -172,6 +316,7 @@ Result>> ArrowMemIndexer::S return std::optional>(); } auto segment = std::make_shared(building_range_.value(), std::move(building_batches_)); + sealed_segments_.push_back(segment); building_batches_.clear(); building_range_.reset(); building_memory_usage_ = 0; @@ -189,12 +334,82 @@ Result>> ArrowMemIndexer::CreateCommitR return readers; } +Result> ArrowMemIndexer::AcquireReadView() { + std::lock_guard lock(mutex_); + if (closed_) { + return Status::Invalid("mem indexer is closed"); + } + 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()); + } + batches.insert(batches.end(), building_batches_.begin(), building_batches_.end()); + return std::shared_ptr(new ReadView(std::move(batches))); +} + +Result>> ArrowMemIndexer::CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_lower_exclusive, + const MemQueryContext& context) { + std::shared_ptr arrow_view = std::dynamic_pointer_cast(view); + if (!arrow_view) { + return Status::Invalid("read view was not created by the Arrow mem indexer"); + } + if (context.read_schema == nullptr || context.read_schema->release == nullptr) { + return Status::Invalid("mem query read schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, + arrow::ImportSchema(context.read_schema)); + // TODO(xinyu.lxy): support batch stats + // Now, the default Arrow indexer has no batch statistics or index metadata yet, so its inner + // query reader intentionally ignores context.predicate instead of using it for candidate + // pruning. + std::vector> readers; + if (arrow_view->GetOffsetRange() && arrow_view->GetOffsetRange()->to > offset_lower_exclusive) { + std::unique_ptr reader = std::make_unique( + arrow_view, offset_lower_exclusive, read_schema, arrow_pool_); + if (context.enable_predicate_filter && context.predicate) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr predicate_reader, + PredicateBatchReader::Create(std::move(reader), context.predicate, memory_pool_)); + reader = std::move(predicate_reader); + } + readers.push_back( + std::make_unique(std::move(reader), memory_pool_)); + } + return readers; +} + +Status ArrowMemIndexer::Reclaim(int64_t committed_offset) { + std::lock_guard lock(mutex_); + if (closed_) { + return Status::Invalid("mem indexer is closed"); + } + // TODO(xinyu.lxy): Consider deferring segment destruction to a reclamation queue. Existing + // read views may pin reclaimed batches, so the last query releasing a view can otherwise pay + // the full buffer destruction cost and observe higher tail latency. + sealed_segments_.erase( + std::remove_if(sealed_segments_.begin(), sealed_segments_.end(), + [committed_offset](const std::shared_ptr& segment) { + return segment->GetOffsetRange().to <= committed_offset; + }), + sealed_segments_.end()); + return Status::OK(); +} + uint64_t ArrowMemIndexer::GetMemoryUsage() const { - return building_memory_usage_; + std::lock_guard lock(mutex_); + uint64_t result = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_segments_) { + result += segment->GetMemoryUsage(); + } + return result; } Status ArrowMemIndexer::Close() { + std::lock_guard lock(mutex_); building_batches_.clear(); + sealed_segments_.clear(); building_memory_usage_ = 0; building_range_.reset(); closed_ = true; diff --git a/src/paimon/core/realtime/arrow_mem_indexer.h b/src/paimon/core/realtime/arrow_mem_indexer.h index 0663cab9..c0eddf96 100644 --- a/src/paimon/core/realtime/arrow_mem_indexer.h +++ b/src/paimon/core/realtime/arrow_mem_indexer.h @@ -18,6 +18,7 @@ #include #include +#include #include #include @@ -30,11 +31,13 @@ class StructArray; } // namespace arrow namespace paimon { +class MemoryPool; /// Internal Arrow-backed implementation of the default `MemIndexer`. class ArrowMemIndexer : public MemIndexer { public: ArrowMemIndexer(const std::shared_ptr& write_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool); Status Write(RealtimeWriteBatch&& write_batch) override; @@ -44,6 +47,14 @@ class ArrowMemIndexer : public MemIndexer { Result>> CreateCommitReaders( const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_lower_exclusive, + const MemQueryContext& context) override; + + Status Reclaim(int64_t committed_offset) override; + uint64_t GetMemoryUsage() const override; Status Close() override; @@ -52,14 +63,21 @@ class ArrowMemIndexer : public MemIndexer { struct StoredBatch { std::shared_ptr data; std::vector row_kinds; + Range offset_range; + uint64_t memory_usage; }; class Segment; + class ReadView; class CommitBatchReader; + class QueryBatchReader; std::shared_ptr write_schema_; + std::shared_ptr memory_pool_; std::shared_ptr arrow_pool_; + mutable std::mutex mutex_; std::vector building_batches_; + std::vector> sealed_segments_; std::optional building_range_; uint64_t building_memory_usage_ = 0; bool closed_ = false; diff --git a/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp b/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp index faa09a57..1d5f8a94 100644 --- a/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp +++ b/src/paimon/core/realtime/arrow_mem_indexer_factory.cpp @@ -36,7 +36,7 @@ Result> ArrowMemIndexerFactory::Create( PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, arrow::ImportSchema(write_schema)); std::shared_ptr arrow_pool = GetArrowPool(memory_pool); - return std::make_shared(imported_schema, arrow_pool); + return std::make_shared(imported_schema, memory_pool, arrow_pool); } } // namespace paimon diff --git a/src/paimon/core/realtime/realtime_context.cpp b/src/paimon/core/realtime/realtime_context.cpp index d1e792ec..80b9bc65 100644 --- a/src/paimon/core/realtime/realtime_context.cpp +++ b/src/paimon/core/realtime/realtime_context.cpp @@ -18,7 +18,9 @@ #include #include +#include #include +#include #include "paimon/arrow/abi.h" #include "paimon/core/realtime/partition_bucket.h" @@ -57,10 +59,51 @@ class RealtimeContext::Impl { return indexer; } + Result> AcquireReadViews() { + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(indexers_.size()); + for (const auto& [partition_bucket, indexer] : indexers_) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, + indexer->AcquireReadView()); + result.push_back(RealtimePartitionBucketView{ + partition_bucket.partition, partition_bucket.bucket, indexer, + std::move(read_view)}); + } + return result; + } + + Status AdvanceCommittedProgress( + int64_t snapshot_id, + const std::vector& committed_offsets) { + if (snapshot_id < 0) { + return Status::Invalid("real-time refresh snapshot id must not be negative"); + } + std::lock_guard lock(mutex_); + if (last_refreshed_snapshot_id_ && snapshot_id < last_refreshed_snapshot_id_.value()) { + return Status::Invalid("real-time committed snapshot cannot move backwards"); + } + if (last_refreshed_snapshot_id_ && snapshot_id == last_refreshed_snapshot_id_.value()) { + return Status::OK(); + } + for (const RealtimePartitionBucketOffset& committed : committed_offsets) { + if (committed.bucket < 0 || committed.offset < 0) { + return Status::Invalid("invalid partition-bucket committed offset"); + } + auto iter = indexers_.find(PartitionBucket(committed.partition, committed.bucket)); + if (iter != indexers_.end()) { + PAIMON_RETURN_NOT_OK(iter->second->Reclaim(committed.offset)); + } + } + last_refreshed_snapshot_id_ = snapshot_id; + return Status::OK(); + } + private: std::shared_ptr factory_; std::mutex mutex_; std::map> indexers_; + std::optional last_refreshed_snapshot_id_; }; Result> RealtimeContext::Create() { @@ -87,4 +130,14 @@ Result> RealtimeContext::GetOrCreateMemIndexer( memory_pool); } +Result> RealtimeContext::AcquireReadViews() { + return impl_->AcquireReadViews(); +} + +Status RealtimeContext::AdvanceCommittedProgress( + int64_t snapshot_id, + const std::vector& committed_offsets) { + return impl_->AdvanceCommittedProgress(snapshot_id, committed_offsets); +} + } // namespace paimon 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 eb6498c1..113f3122 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -20,12 +20,19 @@ #include "paimon/core/table/source/append_only_table_read.h" #include +#include +#include "arrow/c/bridge.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/data_evolution_split_read.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/operation/raw_file_split_read.h" #include "paimon/core/table/source/append_count_reader.h" +#include "paimon/core/table/source/realtime_split.h" +#include "paimon/realtime/mem_indexer.h" #include "paimon/status.h" namespace paimon { @@ -51,6 +58,41 @@ AppendOnlyTableRead::AppendOnlyTableRead(const std::shared_ptr> AppendOnlyTableRead::CreateReader( + const std::shared_ptr& split) { + std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); + if (!realtime_split) { + return CreateDiskReader(split); + } + + std::vector> readers; + readers.reserve(realtime_split->DiskSplits().size() + 1); + for (const std::shared_ptr& disk_split : realtime_split->DiskSplits()) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + CreateDiskReader(disk_split)); + readers.push_back(std::move(disk_reader)); + } + + auto c_read_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*context_->GetReadSchema(), c_read_schema.get())); + ScopeGuard schema_guard([schema = c_read_schema.get()]() { + if (schema->release) { + schema->release(schema); + } + }); + MemQueryContext query_context{c_read_schema.get(), context_->GetPredicate(), + context_->EnablePredicateFilter()}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> memory_readers, + realtime_split->Indexer()->CreateQueryReaders( + realtime_split->ReadView(), realtime_split->CommittedOffset(), query_context)); + for (std::unique_ptr& memory_reader : memory_readers) { + readers.push_back(std::move(memory_reader)); + } + return std::make_unique(std::move(readers), GetMemoryPool()); +} + +Result> AppendOnlyTableRead::CreateDiskReader( const std::shared_ptr& split) { for (const auto& read : split_reads_) { PAIMON_ASSIGN_OR_RAISE(bool matched, read->Match(split, /*force_keep_delete=*/false)); @@ -63,6 +105,12 @@ Result> AppendOnlyTableRead::CreateReader( Result> AppendOnlyTableRead::CreateCountReader( const std::vector>& splits) { + for (const std::shared_ptr& split : splits) { + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented( + "CreateCountReader does not support process-local real-time splits"); + } + } if (context_->GetPredicate() != nullptr) { return Status::NotImplemented( "CreateCountReader with predicate pushdown is not supported yet"); diff --git a/src/paimon/core/table/source/append_only_table_read.h b/src/paimon/core/table/source/append_only_table_read.h index e26b8613..c5d813c8 100644 --- a/src/paimon/core/table/source/append_only_table_read.h +++ b/src/paimon/core/table/source/append_only_table_read.h @@ -51,6 +51,8 @@ class AppendOnlyTableRead : public TableRead { const std::vector>& splits) override; private: + Result> CreateDiskReader(const std::shared_ptr& split); + std::vector> split_reads_; std::shared_ptr context_; }; diff --git a/src/paimon/core/table/source/realtime_split.h b/src/paimon/core/table/source/realtime_split.h new file mode 100644 index 00000000..21f5339f --- /dev/null +++ b/src/paimon/core/table/source/realtime_split.h @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/realtime/mem_indexer.h" +#include "paimon/table/source/split.h" + +namespace paimon { + +/// Process-local split combining committed disk splits and one immutable memory view. +/// +/// This split intentionally has no serialized representation because it carries plugin objects. +class RealtimeSplit : public Split { + public: + RealtimeSplit(std::map partition, int32_t bucket, + std::vector>&& disk_splits, + const std::shared_ptr& indexer, + const std::shared_ptr& read_view, int64_t committed_offset) + : partition_(std::move(partition)), + bucket_(bucket), + disk_splits_(std::move(disk_splits)), + indexer_(indexer), + read_view_(read_view), + committed_offset_(committed_offset) {} + + const std::map& Partition() const { + return partition_; + } + + int32_t Bucket() const { + return bucket_; + } + + const std::vector>& DiskSplits() const { + return disk_splits_; + } + + const std::shared_ptr& Indexer() const { + return indexer_; + } + + const std::shared_ptr& ReadView() const { + return read_view_; + } + + int64_t CommittedOffset() const { + return committed_offset_; + } + + private: + std::map partition_; + int32_t bucket_; + std::vector> disk_splits_; + std::shared_ptr indexer_; + std::shared_ptr read_view_; + int64_t committed_offset_; +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp new file mode 100644 index 00000000..5af891be --- /dev/null +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/core/table/source/realtime_table_scan.h" + +#include +#include +#include +#include + +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" +#include "paimon/core/realtime/partition_bucket.h" +#include "paimon/core/snapshot.h" +#include "paimon/core/table/source/data_split_impl.h" +#include "paimon/core/table/source/plan_impl.h" +#include "paimon/core/table/source/realtime_split.h" +#include "paimon/core/utils/file_store_path_factory.h" +#include "paimon/core/utils/snapshot_manager.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/scan_context.h" +#include "paimon/status.h" + +namespace paimon { + +int64_t RealtimeTableScan::GetCommittedOffset( + const RealtimeSnapshotProperties::OffsetMap& committed_offsets, + const PartitionBucket& partition_bucket) { + auto iter = committed_offsets.find(partition_bucket); + return iter == committed_offsets.end() ? -1 : iter->second; +} + +bool RealtimeTableScan::MatchPartition(const std::map& partition) const { + const auto& filters = scan_filter_->GetPartitionFilters(); + if (filters.empty()) { + return true; + } + for (const auto& filter : filters) { + bool matched = true; + for (const auto& [key, raw_value] : filter) { + auto partition_iter = partition.find(key); + if (partition_iter == partition.end() || partition_iter->second != raw_value) { + matched = false; + break; + } + } + if (matched) { + return true; + } + } + return false; +} + +Result RealtimeTableScan::LoadCommittedOffsets( + const std::shared_ptr& disk_plan) const { + if (!disk_plan->SnapshotId()) { + return RealtimeSnapshotProperties::OffsetMap{}; + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, + snapshot_manager_->LoadSnapshot(disk_plan->SnapshotId().value())); + return RealtimeSnapshotProperties::ReadOffsets(std::optional(std::move(snapshot)), + file_system_); +} + +Result RealtimeTableScan::CollectActiveMemoryViews( + std::vector&& memory_views, + const RealtimeSnapshotProperties::OffsetMap& committed_offsets) const { + MemoryViewMap active_memory; + for (RealtimePartitionBucketView& memory : memory_views) { + if (scan_filter_->GetBucketFilter() && + memory.bucket != scan_filter_->GetBucketFilter().value()) { + continue; + } + if (!MatchPartition(memory.partition)) { + continue; + } + const PartitionBucket key(memory.partition, memory.bucket); + const std::optional memory_range = memory.read_view->GetOffsetRange(); + if (!memory_range || memory_range->to <= GetCommittedOffset(committed_offsets, key)) { + continue; + } + active_memory.emplace(key, std::move(memory)); + } + return active_memory; +} + +Result>> RealtimeTableScan::CreateRealtimeSplits( + const std::vector>& disk_splits, MemoryViewMap&& active_memory, + const RealtimeSnapshotProperties::OffsetMap& committed_offsets) const { + std::map>> disk_by_partition_bucket; + for (const std::shared_ptr& split : disk_splits) { + std::shared_ptr data_split = std::dynamic_pointer_cast(split); + if (!data_split) { + return Status::Invalid("real-time append scan requires process-local data splits"); + } + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + path_factory_->GeneratePartitionVector(data_split->Partition())); + std::map partition(partition_values.begin(), + partition_values.end()); + disk_by_partition_bucket[PartitionBucket(std::move(partition), data_split->Bucket())] + .push_back(split); + } + + // TODO(xinyu.lxy): Support splitting one partition-bucket into multiple real-time splits. + std::vector> result; + for (auto& [key, grouped_disk_splits] : disk_by_partition_bucket) { + auto memory_iter = active_memory.find(key); + if (memory_iter == active_memory.end()) { + result.insert(result.end(), grouped_disk_splits.begin(), grouped_disk_splits.end()); + continue; + } + RealtimePartitionBucketView& memory = memory_iter->second; + result.push_back(std::make_shared( + key.partition, key.bucket, std::move(grouped_disk_splits), memory.indexer, + memory.read_view, GetCommittedOffset(committed_offsets, key))); + active_memory.erase(memory_iter); + } + + for (auto& [key, memory] : active_memory) { + result.push_back(std::make_shared( + key.partition, key.bucket, std::vector>(), memory.indexer, + memory.read_view, GetCommittedOffset(committed_offsets, key))); + } + return result; +} + +RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, + const std::shared_ptr& realtime_context, + const std::shared_ptr& path_factory, + const std::shared_ptr& snapshot_manager, + const std::shared_ptr& file_system, + const std::shared_ptr& scan_filter) + : disk_scan_(std::move(disk_scan)), + realtime_context_(realtime_context), + path_factory_(path_factory), + snapshot_manager_(snapshot_manager), + file_system_(file_system), + scan_filter_(scan_filter) {} + +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(std::shared_ptr disk_plan, disk_scan_->CreatePlan()); + PAIMON_ASSIGN_OR_RAISE(RealtimeSnapshotProperties::OffsetMap committed_offsets, + LoadCommittedOffsets(disk_plan)); + PAIMON_ASSIGN_OR_RAISE(MemoryViewMap active_memory, + CollectActiveMemoryViews(std::move(memory_views), committed_offsets)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> splits, + CreateRealtimeSplits(disk_plan->Splits(), std::move(active_memory), committed_offsets)); + return std::make_shared(disk_plan->SnapshotId(), splits); +} + +} // namespace paimon diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h new file mode 100644 index 00000000..161e832d --- /dev/null +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/core/operation/commit/realtime_snapshot_properties.h" +#include "paimon/core/realtime/partition_bucket.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/result.h" +#include "paimon/table/source/table_scan.h" + +namespace paimon { + +class FileStorePathFactory; +class FileSystem; +class ScanFilter; +class SnapshotManager; + +/// Adds process-local memory splits to a normal append-table batch scan. +class RealtimeTableScan : public TableScan { + public: + RealtimeTableScan(std::unique_ptr&& disk_scan, + const std::shared_ptr& realtime_context, + const std::shared_ptr& path_factory, + const std::shared_ptr& snapshot_manager, + const std::shared_ptr& file_system, + const std::shared_ptr& scan_filter); + + Result> CreatePlan() override; + + private: + using MemoryViewMap = std::map; + + static int64_t GetCommittedOffset( + const RealtimeSnapshotProperties::OffsetMap& committed_offsets, + const PartitionBucket& partition_bucket); + + bool MatchPartition(const std::map& partition) const; + + Result LoadCommittedOffsets( + const std::shared_ptr& disk_plan) const; + + Result CollectActiveMemoryViews( + std::vector&& memory_views, + const RealtimeSnapshotProperties::OffsetMap& committed_offsets) const; + + Result>> CreateRealtimeSplits( + const std::vector>& disk_splits, MemoryViewMap&& active_memory, + const RealtimeSnapshotProperties::OffsetMap& committed_offsets) const; + + std::unique_ptr disk_scan_; + std::shared_ptr realtime_context_; + std::shared_ptr path_factory_; + std::shared_ptr snapshot_manager_; + std::shared_ptr file_system_; + std::shared_ptr scan_filter_; +}; + +} // namespace paimon diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 17add044..62bda904 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -52,6 +52,7 @@ #include "paimon/core/table/source/data_table_stream_scan.h" #include "paimon/core/table/source/merge_tree_split_generator.h" #include "paimon/core/table/source/read_optimized_scan_options.h" +#include "paimon/core/table/source/realtime_table_scan.h" #include "paimon/core/table/source/snapshot/snapshot_reader.h" #include "paimon/core/table/source/split_generator.h" #include "paimon/core/table/system/system_table.h" @@ -61,6 +62,7 @@ #include "paimon/core/utils/index_file_path_factories.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/format/file_format.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/result.h" #include "paimon/scan_context.h" #include "paimon/status.h" @@ -212,6 +214,36 @@ Result> TableScan::Create(std::unique_ptr> NewDataTableScan(const std::shared_ptr& context) { PAIMON_ASSIGN_OR_RAISE( CoreOptions tmp_options, @@ -242,6 +274,8 @@ Result> NewDataTableScan(const std::shared_ptrGetSpecificFileSystem(), {})); core_options.WithCache(context->GetCache()); + + PAIMON_RETURN_NOT_OK(ValidateRealtimeScan(*table_schema, core_options, *context)); // validate options if (core_options.GetBucket() == -1) { if (!table_schema->PrimaryKeys().empty()) { @@ -299,6 +333,12 @@ Result> NewDataTableScan(const std::shared_ptr( /*pk_table=*/pk_table, core_options, snapshot_reader, read_optimized, context->GetLimit()); + if (context->GetRealtimeContext()) { + return std::make_unique( + std::move(batch_scan), context->GetRealtimeContext(), path_factory, + snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(), + context->GetScanFilters()); + } if (!core_options.DataEvolutionEnabled()) { return batch_scan; } diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 2947de9f..ec7e1dff 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -17,7 +17,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -38,11 +40,15 @@ #include "paimon/common/utils/path_util.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_snapshot_properties.h" +#include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/defs.h" #include "paimon/file_store_commit.h" #include "paimon/file_store_write.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" +#include "paimon/realtime/mem_indexer.h" #include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/scan_context.h" @@ -54,12 +60,84 @@ namespace paimon::test { +class ConcurrentTestState { + public: + void WaitForStart() { + ready_threads_.fetch_add(1, std::memory_order_release); + while (!start_.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + } + + void StartWhenReady(int32_t worker_count) { + while (ready_threads_.load(std::memory_order_acquire) < worker_count) { + std::this_thread::yield(); + } + start_.store(true, std::memory_order_release); + } + + void RecordError(const Status& status) { + RecordError(status.ToString()); + } + + void RecordError(std::string error) { + { + std::lock_guard lock(mutex); + errors_.push_back(std::move(error)); + } + stop_.store(true, std::memory_order_release); + progress_cv.notify_all(); + snapshot_cv.notify_all(); + } + + bool RecordErrorIfNotOk(const Status& status) { + if (status.ok()) { + return false; + } + RecordError(status); + return true; + } + + template + bool RecordErrorIfNotOk(const Result& result) { + if (result.ok()) { + return false; + } + RecordError(result.status()); + return true; + } + + bool ShouldStop() const { + return stop_.load(std::memory_order_acquire); + } + + const std::vector& Errors() const { + return errors_; + } + + std::mutex mutex; + std::condition_variable progress_cv; + std::condition_variable snapshot_cv; + + private: + std::atomic start_{false}; + std::atomic stop_{false}; + std::atomic ready_threads_{0}; + std::vector errors_; +}; + class RealtimeWriteInteTest : public ::testing::Test { protected: using Row = std::tuple; using ReadRow = std::tuple; + struct CollectedReadResult { + std::unique_ptr reader; + std::shared_ptr data; + }; + void SetUp() override { + pool_ = GetDefaultPool(); dir_ = UniqueTestDirectory::Create("local"); ASSERT_NE(nullptr, dir_); table_path_ = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); @@ -88,17 +166,27 @@ class RealtimeWriteInteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } - Result> CreateRealtimeWriter() const { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context, - RealtimeContext::Create()); + Result> CreateRealtimeWriter( + const std::shared_ptr& realtime_context) const { WriteContextBuilder builder(table_path_, commit_user_); builder.SetOptions(options_).WithStreamingMode(true).WithRealtimeContext(realtime_context); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); return FileStoreWrite::Create(std::move(context)); } + Result> CreateRealtimeWriter() const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context, + RealtimeContext::Create()); + return CreateRealtimeWriter(realtime_context); + } + Result> MakeBatch(const std::vector& rows, bool partitioned) const { + return MakeBatch(rows, partitioned, /*bucket=*/0); + } + + Result> MakeBatch(const std::vector& rows, bool partitioned, + int32_t bucket) const { if (rows.empty()) { return Status::Invalid("cannot create an empty test batch"); } @@ -125,7 +213,7 @@ class RealtimeWriteInteTest : public ::testing::Test { if (partitioned) { builder.SetPartition({{"pt", partition}}); } - return builder.SetBucket(0).Finish(); + return builder.SetBucket(bucket).Finish(); } static std::vector MakeRows(int64_t first_id, int64_t count, @@ -162,16 +250,30 @@ class RealtimeWriteInteTest : public ::testing::Test { return result; } - Result> ReadRows() const { + Result> CreatePlan( + const std::shared_ptr& realtime_context, + const std::shared_ptr& predicate) const { ScanContextBuilder scan_builder(table_path_); + if (realtime_context) { + scan_builder.WithRealtimeContext(realtime_context); + } + scan_builder.SetPredicate(predicate).WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan_context, scan_builder.SetOptions(options_).Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr scan, TableScan::Create(std::move(scan_context))); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, scan->CreatePlan()); + return scan->CreatePlan(); + } + Result ReadPlan( + const std::shared_ptr& plan, const std::vector& read_fields, + const std::shared_ptr& predicate, bool enable_predicate_filter) const { ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(options_).SetReadFieldNames({"_OFFSET", "id", "payload", "pt"}); + read_builder.SetOptions(options_) + .SetReadFieldNames(read_fields) + .SetPredicate(predicate) + .EnablePredicateFilter(enable_predicate_filter) + .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, TableRead::Create(std::move(read_context))); @@ -179,7 +281,22 @@ class RealtimeWriteInteTest : public ::testing::Test { table_read->CreateReader(plan->Splits())); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, ReadResultCollector::CollectResult(reader.get())); - reader->Close(); + return CollectedReadResult{std::move(reader), std::move(result)}; + } + + Result> ReadRows( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + return ReadRows(plan); + } + + Result> ReadRows(const std::shared_ptr& plan) const { + PAIMON_ASSIGN_OR_RAISE( + CollectedReadResult read_result, + ReadPlan(plan, {"_OFFSET", "id", "payload", "pt"}, /*predicate=*/nullptr, + /*enable_predicate_filter=*/false)); + const std::shared_ptr& result = read_result.data; std::vector rows; if (!result) { @@ -218,6 +335,46 @@ class RealtimeWriteInteTest : public ::testing::Test { return rows; } + Result> ReadRows() const { + return ReadRows(std::shared_ptr()); + } + + Result GetRealtimeMemoryUsage( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::vector views, + realtime_context->AcquireReadViews()); + uint64_t memory_usage = 0; + for (const RealtimePartitionBucketView& view : views) { + memory_usage += view.indexer->GetMemoryUsage(); + } + return memory_usage; + } + + static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { + std::vector seen(static_cast(total_rows), false); + int64_t max_offset = -1; + for (const ReadRow& row : rows) { + const auto& [offset, id, payload, partition] = row; + if (offset < 0 || offset >= total_rows) { + return Status::Invalid("real-time read offset is out of range"); + } + if (seen[static_cast(offset)]) { + return Status::Invalid("real-time read contains duplicate offsets"); + } + if (id != offset || payload != "value-" + std::to_string(id) || partition != "p0") { + return Status::Invalid("real-time read row does not match its offset"); + } + seen[static_cast(offset)] = true; + max_offset = std::max(max_offset, offset); + } + for (int64_t offset = 0; offset <= max_offset; ++offset) { + if (!seen[static_cast(offset)]) { + return Status::Invalid("real-time read contains an offset gap"); + } + } + return Status::OK(); + } + Result ReadCommittedOffsets() const { PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); @@ -233,8 +390,6 @@ class RealtimeWriteInteTest : public ::testing::Test { realtime_commits.insert(realtime_commits.end(), std::make_move_iterator(final_commits.begin()), std::make_move_iterator(final_commits.end())); - // Verify that commit orders prepared offset ranges instead of relying on caller order. - std::reverse(realtime_commits.begin(), realtime_commits.end()); ASSERT_OK(Commit(realtime_commits, prepare_identifier)); ASSERT_OK(writer->Close()); @@ -242,140 +397,750 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_EQ(WithExpectedOffsets(expected_rows), actual_rows); } - void RunConcurrentPrepareTest(int32_t prepare_thread_count) { - CreateTable(/*partition_keys=*/{}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::unique_ptr dir_; + std::string table_path_; + std::string commit_user_ = "realtime_commit_user"; + arrow::FieldVector fields_; + std::shared_ptr schema_; + std::map options_; + std::shared_ptr pool_; +}; - constexpr int64_t kBatchCount = 50; - constexpr int64_t kRowsPerBatch = 4; - std::vector expected_rows = MakeRows( - /*first_id=*/0, kBatchCount * kRowsPerBatch, /*partition=*/"p0"); +TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); +} - std::atomic start{false}; - std::atomic writer_done{false}; - std::atomic ready_threads{0}; - std::atomic next_identifier{0}; - std::mutex result_mutex; - std::vector realtime_commits; - std::vector errors; - std::vector prepare_call_counts(prepare_thread_count, 0); +TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveOffsets) { + options_[Options::TARGET_FILE_ROW_NUM] = "10"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); - auto record_error = [&](const Status& status) { - std::lock_guard lock(result_mutex); - errors.push_back(status.ToString()); - }; - auto append_commits = [&](std::vector&& commits) { - std::lock_guard lock(result_mutex); - realtime_commits.insert(realtime_commits.end(), - std::make_move_iterator(commits.begin()), - std::make_move_iterator(commits.end())); - }; + std::vector expected_rows; + constexpr int64_t kBatchCount = 3; + constexpr int64_t kRowsPerBatch = 10; + for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + std::vector rows = + MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + } - std::thread write_thread([&]() { - ++ready_threads; - while (!start.load(std::memory_order_acquire)) { - std::this_thread::yield(); - } - for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { - std::vector rows = MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, "p0"); - Result> batch_result = - MakeBatch(rows, /*partitioned=*/false); - if (!batch_result.ok()) { - record_error(batch_result.status()); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(Range(0, kBatchCount * kRowsPerBatch - 1), commits[0].offset_range); + std::shared_ptr commit_message = + std::dynamic_pointer_cast(commits[0].commit_message); + ASSERT_NE(nullptr, commit_message); + ASSERT_EQ(3, commit_message->GetNewFilesIncrement().NewFiles().size()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK(writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(WithExpectedOffsets(expected_rows), actual_rows); +} + +TEST_F(RealtimeWriteInteTest, TestCommitOrdersPreparedOffsetRanges) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + + 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(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(Range(0, 2), commits[0].offset_range); + + 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))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_commits.size()); + ASSERT_EQ(Range(3, 4), second_commits[0].offset_range); + + commits.push_back(std::move(second_commits[0])); + std::reverse(commits.begin(), commits.end()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/1)); + ASSERT_OK(writer->Close()); + + std::vector expected_rows = first_rows; + expected_rows.insert(expected_rows.end(), second_rows.begin(), second_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(WithExpectedOffsets(expected_rows), actual_rows); +} + +TEST_F(RealtimeWriteInteTest, TestReadMemoryBeforePrepareCommit) { + 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 rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ(WithExpectedOffsets(rows), actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestReadCommittedDiskAndBuildingMemory) { + 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 disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"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_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, disk_commits.size()); + ASSERT_EQ(Range(0, 2), disk_commits[0].offset_range); + + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector expected_rows = disk_rows; + 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(WithExpectedOffsets(expected_rows), actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { + 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::shared_ptr scan_predicate = + PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(1))); + std::shared_ptr read_predicate = + PredicateBuilder::GreaterThan(/*field_index=*/1, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(1))); + const std::vector read_fields = {"payload", "id"}; + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()), + arrow::field("id", arrow::int64())}); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"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::shared_ptr memory_plan, + CreatePlan(realtime_context, scan_predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult memory_result, + ReadPlan(memory_plan, read_fields, read_predicate, /*enable_predicate_filter=*/true)); + std::shared_ptr expected_memory = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "value-2", 2] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, memory_result.data); + ASSERT_TRUE( + std::make_shared(expected_memory)->Equals(*memory_result.data)); + + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr union_plan, + CreatePlan(realtime_context, scan_predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult union_result, + ReadPlan(union_plan, read_fields, read_predicate, /*enable_predicate_filter=*/true)); + std::shared_ptr expected_union = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "value-2", 2], + [0, "value-3", 3], + [0, "value-4", 4], + [0, "value-5", 5] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, union_result.data); + ASSERT_TRUE(std::make_shared(expected_union)->Equals(*union_result.data)); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { + options_[Options::FILE_FORMAT] = "parquet"; + options_[Options::WRITE_BATCH_SIZE] = "1"; + options_["parquet.page.size"] = "1"; + options_["parquet.enable-dictionary"] = "false"; + options_["parquet.write.enable-page-index"] = "true"; + options_["parquet.read.enable-page-index-filter"] = "true"; + 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::shared_ptr predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"id", FieldType::BIGINT, + Literal(static_cast(1))); + + std::vector disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"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_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult result, + ReadPlan(plan, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/false)); + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, 1, "value-1", "p0"], + [0, 3, "value-3", "p0"], + [0, 4, "value-4", "p0"], + [0, 5, "value-5", "p0"] + ])") + .ValueOrDie(); + ASSERT_NE(nullptr, result.data); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestRefreshCommittedSnapshotReclaimsMemory) { + 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 disk_rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"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_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, disk_commits.size()); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector memory_rows = MakeRows(/*first_id=*/10, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + std::vector expected_rows = disk_rows; + expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector read1, ReadRows(realtime_context)); + ASSERT_EQ(WithExpectedOffsets(expected_rows), read1); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh, + GetRealtimeMemoryUsage(realtime_context)); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(latest_snapshot.has_value()); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot->Id())); + + ASSERT_OK_AND_ASSIGN(std::vector read2, ReadRows(realtime_context)); + ASSERT_EQ(read1, read2); + 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()); +} + +TEST_F(RealtimeWriteInteTest, TestPlanPinsMemoryAcrossRefresh) { + 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 disk_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"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_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector memory_rows = MakeRows(/*first_id=*/3, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch(memory_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(latest_snapshot.has_value()); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot->Id())); + + std::vector expected_rows = disk_rows; + expected_rows.insert(expected_rows.end(), memory_rows.begin(), memory_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector pinned_rows, ReadRows(pinned_plan)); + ASSERT_EQ(WithExpectedOffsets(expected_rows), pinned_rows); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr refreshed_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(std::vector refreshed_rows, ReadRows(refreshed_plan)); + ASSERT_EQ(pinned_rows, refreshed_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestRepeatedCommitReadAndRefresh) { + 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(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + + constexpr int64_t kRoundCount = 3; + constexpr int64_t kRowsPerRound = 4; + std::vector expected_rows; + for (int64_t round = 0; round < kRoundCount; ++round) { + std::vector rows = MakeRows(round * kRowsPerRound, kRowsPerRound, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/round)); + ASSERT_EQ(1, commits.size()); + ASSERT_EQ(Range(round * kRowsPerRound, (round + 1) * kRowsPerRound - 1), + commits[0].offset_range); + ASSERT_OK(Commit(commits, /*commit_identifier=*/round)); + expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + + ASSERT_OK_AND_ASSIGN(std::vector read_before_refresh, ReadRows(realtime_context)); + ASSERT_EQ(WithExpectedOffsets(expected_rows), read_before_refresh); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_before_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_GT(memory_usage_before_refresh, 0); + + ASSERT_OK_AND_ASSIGN(std::optional latest_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(latest_snapshot.has_value()); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot->Id())); + + ASSERT_OK_AND_ASSIGN(std::vector read_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(read_before_refresh, read_after_refresh); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage_after_refresh, + GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage_after_refresh); + } + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { + 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(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + + constexpr int32_t kPrepareThreadCount = 4; + constexpr int32_t kReadThreadCount = 4; + constexpr int64_t kBatchCount = 12; + constexpr int64_t kRowsPerBatch = 2; + constexpr int64_t kTotalRows = kBatchCount * kRowsPerBatch; + + std::atomic writer_done{false}; + std::atomic prepare_done{false}; + std::atomic commit_done{false}; + std::atomic refresh_done{false}; + std::atomic next_prepare_identifier{0}; + std::atomic commit_count{0}; + std::atomic refresh_count{0}; + ConcurrentTestState state; + std::map pending_commits; + std::deque pending_snapshot_ids; + std::vector prepare_call_counts(kPrepareThreadCount, 0); + std::vector read_call_counts(kReadThreadCount, 0); + + auto enqueue_prepared_commits = [&](std::vector&& commits) { + std::string error; + { + std::lock_guard lock(state.mutex); + for (RealtimeCommitProgress& commit : commits) { + int64_t offset_from = commit.offset_range.from; + if (!pending_commits.emplace(offset_from, std::move(commit)).second) { + error = "duplicate prepared real-time offset range"; break; } - Status status = writer->Write(std::move(batch_result).value()); - if (!status.ok()) { - record_error(status); + } + } + if (!error.empty()) { + state.RecordError(error); + } + state.progress_cv.notify_all(); + }; + + std::thread write_thread([&]() { + state.WaitForStart(); + for (int64_t batch_index = 0; + batch_index < kBatchCount && !state.ShouldStop(); ++batch_index) { + std::vector rows = + MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); + Result> batch_result = + MakeBatch(rows, /*partitioned=*/false); + if (state.RecordErrorIfNotOk(batch_result)) { + break; + } + Status status = writer->Write(std::move(batch_result).value()); + if (state.RecordErrorIfNotOk(status)) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + // Pause midway until one refresh completes to guarantee write and refresh overlap. + if (batch_index + 1 == kBatchCount / 2) { + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + while (refresh_count.load(std::memory_order_acquire) == 0 && + !state.ShouldStop() && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + if (refresh_count.load(std::memory_order_acquire) == 0 && + !state.ShouldStop()) { + state.RecordError("timed out waiting for a refresh while writing"); break; } - std::this_thread::sleep_for(std::chrono::microseconds(50)); } - writer_done.store(true, std::memory_order_release); + } + writer_done.store(true, std::memory_order_release); + }); + + std::vector prepare_threads; + prepare_threads.reserve(kPrepareThreadCount); + for (int32_t thread_index = 0; thread_index < kPrepareThreadCount; ++thread_index) { + prepare_threads.emplace_back([&, thread_index]() { + state.WaitForStart(); + do { + int64_t identifier = next_prepare_identifier.fetch_add(1); + Result> result = + writer->PrepareCommitWithProgress(identifier); + ++prepare_call_counts[thread_index]; + if (state.RecordErrorIfNotOk(result)) { + break; + } + enqueue_prepared_commits(std::move(result).value()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } while (!writer_done.load(std::memory_order_acquire) && + !state.ShouldStop()); }); + } - std::vector prepare_threads; - prepare_threads.reserve(prepare_thread_count); - for (int32_t thread_index = 0; thread_index < prepare_thread_count; ++thread_index) { - prepare_threads.emplace_back([&, thread_index]() { - ++ready_threads; - while (!start.load(std::memory_order_acquire)) { - std::this_thread::yield(); + std::thread commit_thread([&]() { + state.WaitForStart(); + int64_t next_offset = 0; + int64_t commit_identifier = 0; + while (!state.ShouldStop()) { + std::optional next_commit; + { + std::unique_lock lock(state.mutex); + state.progress_cv.wait(lock, [&]() { + return state.ShouldStop() || + pending_commits.count(next_offset) > 0 || + prepare_done.load(std::memory_order_acquire); + }); + if (state.ShouldStop()) { + break; } - do { - int64_t identifier = next_identifier.fetch_add(1); - Result> result = - writer->PrepareCommitWithProgress(identifier); - ++prepare_call_counts[thread_index]; - if (!result.ok()) { - record_error(result.status()); + auto iter = pending_commits.find(next_offset); + if (iter == pending_commits.end()) { + if (prepare_done.load(std::memory_order_acquire)) { + if (!pending_commits.empty()) { + lock.unlock(); + state.RecordError("prepared real-time offset ranges contain a gap"); + } break; } - append_commits(std::move(result).value()); - std::this_thread::sleep_for(std::chrono::microseconds(100)); - } while (!writer_done.load(std::memory_order_acquire)); - }); - } + continue; + } + next_commit = std::move(iter->second); + pending_commits.erase(iter); + } - while (ready_threads.load(std::memory_order_acquire) < prepare_thread_count + 1) { - std::this_thread::yield(); + std::vector commits; + commits.push_back(std::move(next_commit).value()); + int64_t committed_offset = commits[0].offset_range.to; + Status status = Commit(commits, commit_identifier++); + if (state.RecordErrorIfNotOk(status)) { + break; + } + next_offset = committed_offset + 1; + Result> snapshot_result = snapshot_manager.LatestSnapshot(); + if (state.RecordErrorIfNotOk(snapshot_result)) { + break; + } + std::optional snapshot = std::move(snapshot_result).value(); + if (!snapshot) { + state.RecordError("real-time commit did not create a snapshot"); + break; + } + { + std::lock_guard lock(state.mutex); + pending_snapshot_ids.push_back(snapshot->Id()); + } + ++commit_count; + state.snapshot_cv.notify_all(); } - start.store(true, std::memory_order_release); + commit_done.store(true, std::memory_order_release); + state.snapshot_cv.notify_all(); + }); - write_thread.join(); - for (std::thread& thread : prepare_threads) { - thread.join(); + std::thread refresh_thread([&]() { + state.WaitForStart(); + while (!state.ShouldStop()) { + std::optional snapshot_id; + { + std::unique_lock lock(state.mutex); + state.snapshot_cv.wait(lock, [&]() { + return state.ShouldStop() || !pending_snapshot_ids.empty() || + commit_done.load(std::memory_order_acquire); + }); + if (state.ShouldStop()) { + break; + } + if (pending_snapshot_ids.empty()) { + if (commit_done.load(std::memory_order_acquire)) { + break; + } + continue; + } + snapshot_id = pending_snapshot_ids.front(); + pending_snapshot_ids.pop_front(); + } + Status status = writer->RefreshCommittedSnapshot(snapshot_id.value()); + if (state.RecordErrorIfNotOk(status)) { + break; + } + ++refresh_count; } + refresh_done.store(true, std::memory_order_release); + }); + + std::vector read_threads; + read_threads.reserve(kReadThreadCount); + for (int32_t thread_index = 0; thread_index < kReadThreadCount; ++thread_index) { + read_threads.emplace_back([&, thread_index]() { + state.WaitForStart(); + while (!refresh_done.load(std::memory_order_acquire) && + !state.ShouldStop()) { + Result> result = ReadRows(realtime_context); + ++read_call_counts[thread_index]; + if (state.RecordErrorIfNotOk(result)) { + break; + } + Status status = ValidateReadPrefix(result.value(), kTotalRows); + if (state.RecordErrorIfNotOk(status)) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + }); + } - ASSERT_TRUE(errors.empty()) << (errors.empty() ? "" : errors.front()); - for (int32_t call_count : prepare_call_counts) { - ASSERT_GT(call_count, 0); + constexpr int32_t kWorkerCount = 1 + kPrepareThreadCount + 1 + 1 + kReadThreadCount; + state.StartWhenReady(kWorkerCount); + + write_thread.join(); + for (std::thread& thread : prepare_threads) { + thread.join(); + } + if (!state.ShouldStop()) { + Result> final_result = + writer->PrepareCommitWithProgress(next_prepare_identifier.fetch_add(1)); + if (!state.RecordErrorIfNotOk(final_result)) { + enqueue_prepared_commits(std::move(final_result).value()); } - FinalizeCommitAndCheck(writer.get(), std::move(realtime_commits), - next_identifier.fetch_add(1), std::move(expected_rows)); } + prepare_done.store(true, std::memory_order_release); + state.progress_cv.notify_all(); - std::unique_ptr dir_; - std::string table_path_; - std::string commit_user_ = "realtime_commit_user"; - arrow::FieldVector fields_; - std::shared_ptr schema_; - std::map options_; -}; + commit_thread.join(); + refresh_thread.join(); + for (std::thread& thread : read_threads) { + thread.join(); + } -TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { - CreateTable(/*partition_keys=*/{}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); - std::vector rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch(rows, /*partitioned=*/false)); - ASSERT_OK(writer->Write(std::move(batch))); - FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); + ASSERT_TRUE(state.Errors().empty()) + << (state.Errors().empty() ? "" : state.Errors().front()); + for (int32_t call_count : prepare_call_counts) { + ASSERT_GT(call_count, 0); + } + for (int32_t call_count : read_call_counts) { + ASSERT_GT(call_count, 0); + } + ASSERT_GE(commit_count.load(), 2); + ASSERT_GE(refresh_count.load(), 2); + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows(realtime_context)); + ASSERT_EQ(kTotalRows, static_cast(final_rows.size())); + ASSERT_OK(ValidateReadPrefix(final_rows, kTotalRows)); + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap committed_offsets, + ReadCommittedOffsets()); + ASSERT_EQ(kTotalRows - 1, + committed_offsets.at(PartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + ASSERT_OK(writer->Close()); } TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); - std::vector expected_rows; - for (int64_t partition_index = 0; partition_index < 3; ++partition_index) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector> disk_rows; + for (int64_t partition_index = 0; partition_index < 2; ++partition_index) { std::string partition = "p" + std::to_string(partition_index); std::vector rows = MakeRows(partition_index * 10, /*count=*/10, partition); ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeBatch(rows, /*partitioned=*/true)); ASSERT_OK(writer->Write(std::move(batch))); - expected_rows.insert(expected_rows.end(), rows.begin(), rows.end()); + disk_rows.push_back(std::move(rows)); } - FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, - std::move(expected_rows)); + ASSERT_OK_AND_ASSIGN(std::vector disk_commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, disk_commits.size()); + ASSERT_OK(Commit(disk_commits, /*commit_identifier=*/0)); + + std::vector p0_memory_rows = MakeRows(/*first_id=*/20, /*count=*/5, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_memory_batch, + MakeBatch(p0_memory_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p0_memory_batch))); + std::vector p2_memory_rows = MakeRows(/*first_id=*/30, /*count=*/5, /*partition=*/"p2"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p2_memory_batch, + MakeBatch(p2_memory_rows, /*partitioned=*/true)); + ASSERT_OK(writer->Write(std::move(p2_memory_batch))); + + // p0 has disk and memory rows, p1 is disk-only, and p2 is memory-only. + std::vector expected_rows = disk_rows[0]; + expected_rows.insert(expected_rows.end(), p0_memory_rows.begin(), p0_memory_rows.end()); + expected_rows.insert(expected_rows.end(), disk_rows[1].begin(), disk_rows[1].end()); + expected_rows.insert(expected_rows.end(), p2_memory_rows.begin(), p2_memory_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ(WithExpectedOffsets(expected_rows), actual_rows); + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap committed_offsets, ReadCommittedOffsets()); - ASSERT_EQ(3, committed_offsets.size()); - for (int64_t partition_index = 0; partition_index < 3; ++partition_index) { - std::map partition = { - {"pt", "p" + std::to_string(partition_index)}}; - PartitionBucket partition_bucket(partition, /*bucket=*/0); + ASSERT_EQ(2, committed_offsets.size()); + for (int64_t partition_index = 0; partition_index < 2; ++partition_index) { + PartitionBucket partition_bucket({{"pt", "p" + std::to_string(partition_index)}}, + /*bucket=*/0); ASSERT_EQ(9, committed_offsets.at(partition_bucket)); } + ASSERT_EQ(committed_offsets.end(), + committed_offsets.find(PartitionBucket({{"pt", "p2"}}, /*bucket=*/0))); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestMultipleBucketsRestoreIndependentOffsets) { + options_[Options::BUCKET] = "2"; + CreateTable(/*partition_keys=*/{}); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, CreateRealtimeWriter()); + std::vector bucket0_disk_rows = + MakeRows(/*first_id=*/0, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket0_disk_batch, + MakeBatch(bucket0_disk_rows, /*partitioned=*/false, /*bucket=*/0)); + ASSERT_OK(first_writer->Write(std::move(bucket0_disk_batch))); + std::vector bucket1_disk_rows = + MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket1_disk_batch, + MakeBatch(bucket1_disk_rows, /*partitioned=*/false, /*bucket=*/1)); + ASSERT_OK(first_writer->Write(std::move(bucket1_disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_commits, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, first_commits.size()); + ASSERT_OK(Commit(first_commits, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(realtime_context)); + std::vector bucket0_memory_rows = + MakeRows(/*first_id=*/2, /*count=*/1, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket0_memory_batch, + MakeBatch(bucket0_memory_rows, /*partitioned=*/false, /*bucket=*/0)); + ASSERT_OK(second_writer->Write(std::move(bucket0_memory_batch))); + std::vector bucket1_memory_rows = + MakeRows(/*first_id=*/13, /*count=*/1, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket1_memory_batch, + MakeBatch(bucket1_memory_rows, /*partitioned=*/false, /*bucket=*/1)); + ASSERT_OK(second_writer->Write(std::move(bucket1_memory_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_commits, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, second_commits.size()); + std::map prepared_ranges; + for (const RealtimeCommitProgress& commit : second_commits) { + prepared_ranges.emplace(commit.bucket, commit.offset_range); + } + ASSERT_EQ(Range(2, 2), prepared_ranges.at(0)); + ASSERT_EQ(Range(3, 3), prepared_ranges.at(1)); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + std::map expected_offsets = { + {0, 0}, {1, 1}, {2, 2}, {10, 0}, {11, 1}, {12, 2}, {13, 3}}; + ASSERT_EQ(expected_offsets.size(), actual_rows.size()); + for (const ReadRow& row : actual_rows) { + const auto& [offset, id, payload, partition] = row; + auto iter = expected_offsets.find(id); + ASSERT_NE(expected_offsets.end(), iter); + ASSERT_EQ(iter->second, offset); + ASSERT_EQ("value-" + std::to_string(id), payload); + ASSERT_EQ("p0", partition); + expected_offsets.erase(iter); + } + ASSERT_TRUE(expected_offsets.empty()); + + ASSERT_OK(Commit(second_commits, /*commit_identifier=*/1)); + ASSERT_OK(second_writer->Close()); + ASSERT_OK_AND_ASSIGN(RealtimeSnapshotProperties::OffsetMap committed_offsets, + ReadCommittedOffsets()); + ASSERT_EQ(2, committed_offsets.size()); + ASSERT_EQ(2, committed_offsets.at(PartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_EQ(3, committed_offsets.at(PartitionBucket(/*partition=*/{}, /*bucket=*/1))); } TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { @@ -413,12 +1178,4 @@ TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { ASSERT_EQ(4, second_committed_offsets.at(partition_bucket)); } -TEST_F(RealtimeWriteInteTest, TestConcurrentWriteAndPrepareCommit) { - RunConcurrentPrepareTest(/*prepare_thread_count=*/1); -} - -TEST_F(RealtimeWriteInteTest, TestConcurrentWriteAndMultiplePrepareCommitThreads) { - RunConcurrentPrepareTest(/*prepare_thread_count=*/4); -} - } // namespace paimon::test From 97bfcebd9bc78ba80ba0daf4488636043672dfca Mon Sep 17 00:00:00 2001 From: lxy264173 Date: Sun, 2 Aug 2026 12:35:09 +0800 Subject: [PATCH 4/4] fix clang-tidy --- src/paimon/core/realtime/realtime_context.cpp | 12 ++--- test/inte/realtime_write_inte_test.cpp | 46 ++++++++----------- 2 files changed, 24 insertions(+), 34 deletions(-) diff --git a/src/paimon/core/realtime/realtime_context.cpp b/src/paimon/core/realtime/realtime_context.cpp index 80b9bc65..1fd5a350 100644 --- a/src/paimon/core/realtime/realtime_context.cpp +++ b/src/paimon/core/realtime/realtime_context.cpp @@ -66,16 +66,15 @@ class RealtimeContext::Impl { for (const auto& [partition_bucket, indexer] : indexers_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, indexer->AcquireReadView()); - result.push_back(RealtimePartitionBucketView{ - partition_bucket.partition, partition_bucket.bucket, indexer, - std::move(read_view)}); + result.push_back(RealtimePartitionBucketView{partition_bucket.partition, + partition_bucket.bucket, indexer, + std::move(read_view)}); } return result; } Status AdvanceCommittedProgress( - int64_t snapshot_id, - const std::vector& committed_offsets) { + int64_t snapshot_id, const std::vector& committed_offsets) { if (snapshot_id < 0) { return Status::Invalid("real-time refresh snapshot id must not be negative"); } @@ -135,8 +134,7 @@ Result> RealtimeContext::AcquireReadVie } Status RealtimeContext::AdvanceCommittedProgress( - int64_t snapshot_id, - const std::vector& committed_offsets) { + int64_t snapshot_id, const std::vector& committed_offsets) { return impl_->AdvanceCommittedProgress(snapshot_id, committed_offsets); } diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index ec7e1dff..542125d8 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -265,9 +265,10 @@ class RealtimeWriteInteTest : public ::testing::Test { return scan->CreatePlan(); } - Result ReadPlan( - const std::shared_ptr& plan, const std::vector& read_fields, - const std::shared_ptr& predicate, bool enable_predicate_filter) const { + Result ReadPlan(const std::shared_ptr& plan, + const std::vector& read_fields, + const std::shared_ptr& predicate, + bool enable_predicate_filter) const { ReadContextBuilder read_builder(table_path_); read_builder.SetOptions(options_) .SetReadFieldNames(read_fields) @@ -618,10 +619,9 @@ TEST_F(RealtimeWriteInteTest, TestDiskPredicatePushdownWithoutMemoryFiltering) { ASSERT_OK(writer->Write(std::move(memory_batch))); ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); - ASSERT_OK_AND_ASSIGN( - CollectedReadResult result, - ReadPlan(plan, {"id", "payload", "pt"}, predicate, - /*enable_predicate_filter=*/false)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/false)); std::shared_ptr result_type = arrow::struct_( {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::utf8())}); @@ -816,8 +816,8 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { std::thread write_thread([&]() { state.WaitForStart(); - for (int64_t batch_index = 0; - batch_index < kBatchCount && !state.ShouldStop(); ++batch_index) { + for (int64_t batch_index = 0; batch_index < kBatchCount && !state.ShouldStop(); + ++batch_index) { std::vector rows = MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, /*partition=*/"p0"); Result> batch_result = @@ -833,13 +833,11 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { // Pause midway until one refresh completes to guarantee write and refresh overlap. if (batch_index + 1 == kBatchCount / 2) { auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); - while (refresh_count.load(std::memory_order_acquire) == 0 && - !state.ShouldStop() && + while (refresh_count.load(std::memory_order_acquire) == 0 && !state.ShouldStop() && std::chrono::steady_clock::now() < deadline) { std::this_thread::sleep_for(std::chrono::milliseconds(1)); } - if (refresh_count.load(std::memory_order_acquire) == 0 && - !state.ShouldStop()) { + if (refresh_count.load(std::memory_order_acquire) == 0 && !state.ShouldStop()) { state.RecordError("timed out waiting for a refresh while writing"); break; } @@ -863,8 +861,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { } enqueue_prepared_commits(std::move(result).value()); std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } while (!writer_done.load(std::memory_order_acquire) && - !state.ShouldStop()); + } while (!writer_done.load(std::memory_order_acquire) && !state.ShouldStop()); }); } @@ -877,8 +874,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { { std::unique_lock lock(state.mutex); state.progress_cv.wait(lock, [&]() { - return state.ShouldStop() || - pending_commits.count(next_offset) > 0 || + return state.ShouldStop() || pending_commits.count(next_offset) > 0 || prepare_done.load(std::memory_order_acquire); }); if (state.ShouldStop()) { @@ -963,8 +959,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { for (int32_t thread_index = 0; thread_index < kReadThreadCount; ++thread_index) { read_threads.emplace_back([&, thread_index]() { state.WaitForStart(); - while (!refresh_done.load(std::memory_order_acquire) && - !state.ShouldStop()) { + while (!refresh_done.load(std::memory_order_acquire) && !state.ShouldStop()) { Result> result = ReadRows(realtime_context); ++read_call_counts[thread_index]; if (state.RecordErrorIfNotOk(result)) { @@ -1002,8 +997,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { thread.join(); } - ASSERT_TRUE(state.Errors().empty()) - << (state.Errors().empty() ? "" : state.Errors().front()); + ASSERT_TRUE(state.Errors().empty()) << (state.Errors().empty() ? "" : state.Errors().front()); for (int32_t call_count : prepare_call_counts) { ASSERT_GT(call_count, 0); } @@ -1079,13 +1073,11 @@ TEST_F(RealtimeWriteInteTest, TestMultipleBucketsRestoreIndependentOffsets) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, CreateRealtimeWriter()); - std::vector bucket0_disk_rows = - MakeRows(/*first_id=*/0, /*count=*/2, /*partition=*/"p0"); + std::vector bucket0_disk_rows = MakeRows(/*first_id=*/0, /*count=*/2, /*partition=*/"p0"); ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket0_disk_batch, MakeBatch(bucket0_disk_rows, /*partitioned=*/false, /*bucket=*/0)); ASSERT_OK(first_writer->Write(std::move(bucket0_disk_batch))); - std::vector bucket1_disk_rows = - MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p0"); + std::vector bucket1_disk_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p0"); ASSERT_OK_AND_ASSIGN(std::unique_ptr bucket1_disk_batch, MakeBatch(bucket1_disk_rows, /*partitioned=*/false, /*bucket=*/1)); ASSERT_OK(first_writer->Write(std::move(bucket1_disk_batch))); @@ -1120,8 +1112,8 @@ TEST_F(RealtimeWriteInteTest, TestMultipleBucketsRestoreIndependentOffsets) { ASSERT_EQ(Range(3, 3), prepared_ranges.at(1)); ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); - std::map expected_offsets = { - {0, 0}, {1, 1}, {2, 2}, {10, 0}, {11, 1}, {12, 2}, {13, 3}}; + std::map expected_offsets = {{0, 0}, {1, 1}, {2, 2}, {10, 0}, + {11, 1}, {12, 2}, {13, 3}}; ASSERT_EQ(expected_offsets.size(), actual_rows.size()); for (const ReadRow& row : actual_rows) { const auto& [offset, id, payload, partition] = row;