diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 4d65743ab..b3fa630de 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -27,10 +27,7 @@ namespace paimon { class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: /// Creates an Arrow-backed store for one partition and bucket. - Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) override; + Result> Create(RealtimeStoreCreateRequest&& request) override; }; } // namespace paimon diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index d02952acd..6ae81f1f4 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -27,6 +27,7 @@ #include #include +#include "arrow/c/abi.h" #include "paimon/reader/batch_reader.h" #include "paimon/realtime/offset_range.h" #include "paimon/record_batch.h" @@ -41,10 +42,33 @@ namespace paimon { class MemoryPool; class Predicate; -/// A table record batch and its framework-assigned contiguous offset range. +enum class PAIMON_EXPORT RealtimeStoreMode { + APPEND_ONLY, + PRIMARY_KEY, +}; + +/// Parameters used by a `RealtimeStoreFactory` to create a store. +struct PAIMON_EXPORT RealtimeStoreCreateRequest { + /// Schema whose ownership is transferred to the factory. Append mode receives the complete + /// table write schema. Primary-key mode receives the realtime primary-key transport schema: + /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields]. + std::unique_ptr<::ArrowSchema> write_schema; + /// Table options available to the store implementation. + std::map options; + /// Memory pool for allocations retained by the store. + std::shared_ptr memory_pool; + /// Table mode implemented by the store. + RealtimeStoreMode mode = RealtimeStoreMode::APPEND_ONLY; + /// Statistics collected by append-only stores. + StatisticsMode statistics_mode = StatisticsMode::NONE; +}; + +/// A record batch and its framework-assigned contiguous offset range. /// -/// The batch contains only table write fields. Row `i` is associated with -/// `offset_range.begin + i`; the offset is progress metadata and is not a table field. +/// Append-mode batches contain table write fields, and row `i` has offset +/// `offset_range.begin + i`. Primary-key batches use the realtime primary-key transport schema, +/// are sorted by full primary key then sequence number, and retain the original offset in +/// `_REALTIME_OFFSET`. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -79,7 +103,11 @@ class PAIMON_EXPORT RealtimeReadView { /// Parameters used by a `RealtimeStore` to create readers for a query. struct PAIMON_EXPORT RealtimeQueryContext { - /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added. + /// Append mode receives the requested output fields before the mandatory leading + /// `_VALUE_KIND` field is added. Primary-key mode receives the requested realtime primary-key + /// transport schema. + /// This schema is borrowed and remains valid only during `CreateQueryReaders`; plugins must + /// import or copy it synchronously. ::ArrowSchema* read_schema; /// Predicate using field indexes from `read_schema`. std::shared_ptr predicate; @@ -116,9 +144,10 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// - /// Concatenating the returned readers must produce every sealed row exactly once and in write - /// order. Each output batch contains `_VALUE_KIND` followed by all fields from the factory's - /// `write_schema`. + /// The returned readers collectively expose every sealed row exactly once. Append-mode readers + /// preserve write order and contain `_VALUE_KIND` followed by table write fields. Primary-key + /// readers use the realtime primary-key transport schema; each reader's complete stream is + /// sorted by full primary key then sequence number. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -128,13 +157,16 @@ class PAIMON_EXPORT RealtimeStore { /// also provide a consistent snapshot when a write or seal is in progress. virtual Result> AcquireReadView() = 0; - /// Creates readers over rows in `view` whose offsets are greater than or equal to - /// `offset_begin`. + /// Creates readers over rows in `view`. Append mode returns rows whose offsets are greater than + /// or equal to `offset_begin`; primary-key mode ignores `offset_begin`. /// - /// Each output batch contains `_VALUE_KIND` first, followed by the fields requested by - /// `context.read_schema` except a duplicate `_VALUE_KIND`. Concatenating all returned readers - /// must produce every matching row once. Paimon retains `view` for the lifetime of the - /// resulting framework reader. + /// Append-mode batches contain `_VALUE_KIND` followed by the requested fields except a + /// duplicate `_VALUE_KIND`, and collectively expose every matching row exactly once. + /// Primary-key batches use the requested realtime primary-key transport schema, including + /// nested field-ID alignment, and may contain multiple mutations per key; each reader's + /// complete stream is sorted by full primary key then sequence number, and the readers + /// collectively expose every raw mutation exactly once. Paimon retains `view` for the lifetime + /// of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; @@ -158,15 +190,9 @@ class PAIMON_EXPORT RealtimeStoreFactory { virtual ~RealtimeStoreFactory() = default; /// Creates a store configured with the supplied schema, statistics, options, and memory pool. - /// @param write_schema Complete table write schema whose ownership is transferred to the - /// factory. The factory may consume it or retain it in the created store. - /// @param statistics_mode Framework-parsed statistics collection mode. - /// @param options Effective table options available to the store. - /// @param memory_pool Memory pool provided by the write context. - virtual Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) = 0; + /// Creates a store for the requested table mode. + /// The factory consumes `request`, including ownership of `request.write_schema`. + virtual Result> Create(RealtimeStoreCreateRequest&& request) = 0; }; } // namespace paimon diff --git a/include/paimon/utils/special_field_ids.h b/include/paimon/utils/special_field_ids.h index 829f29889..5219d72db 100644 --- a/include/paimon/utils/special_field_ids.h +++ b/include/paimon/utils/special_field_ids.h @@ -42,6 +42,8 @@ class SpecialFieldIds { /// Special field ID reserved for index score. Value: CPP_FIELD_ID_END - 1 inline static constexpr int32_t INDEX_SCORE = CPP_FIELD_ID_END - 1; + /// Special field ID reserved for realtime offset. Value: CPP_FIELD_ID_END - 2 + inline static constexpr int32_t REALTIME_OFFSET = CPP_FIELD_ID_END - 2; /// Lowest field ID reserved for system fields; IDs at or above it are excluded from the /// highest field ID of a schema. Value: INT32_MAX / 2 diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 3de2b667e..eb3579e29 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -384,9 +384,12 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/realtime_primary_key_reader.cpp + core/realtime/primary_key_realtime_store.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp + core/realtime/realtime_primary_key_writer.cpp core/postpone/postpone_bucket_writer.cpp core/schema/arrow_schema_validator.cpp core/schema/schema_manager.cpp @@ -790,6 +793,8 @@ if(PAIMON_BUILD_TESTS) core/manifest/index_manifest_file_handler_test.cpp core/memory/writer_memory_manager_test.cpp core/realtime/arrow_realtime_store_test.cpp + core/realtime/primary_key_realtime_store_test.cpp + core/realtime/realtime_primary_key_reader_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp diff --git a/src/paimon/common/table/special_fields.h b/src/paimon/common/table/special_fields.h index 74b95b19c..8908f9f0c 100644 --- a/src/paimon/common/table/special_fields.h +++ b/src/paimon/common/table/special_fields.h @@ -21,6 +21,7 @@ #include #include #include +#include #include "arrow/type_fwd.h" #include "paimon/common/types/data_field.h" @@ -66,13 +67,20 @@ struct SpecialFields { return data_field; } + static const DataField& RealtimeOffset() { + static const DataField data_field = + DataField(SpecialFieldIds::REALTIME_OFFSET, + arrow::field("_REALTIME_OFFSET", arrow::int64(), false)); + return data_field; + } + static bool IsSystemField(const std::string& field_name) { if (StringUtils::StartsWith(field_name, KEY_FIELD_PREFIX)) { return true; } return field_name == SequenceNumber().Name() || field_name == ValueKind().Name() || field_name == RowKind().Name() || field_name == RowId().Name() || - field_name == IndexScore().Name(); + field_name == IndexScore().Name() || field_name == RealtimeOffset().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 68e805fd6..58a025ba2 100644 --- a/src/paimon/common/table/special_fields_test.cpp +++ b/src/paimon/common/table/special_fields_test.cpp @@ -55,6 +55,13 @@ TEST(SpecialFieldsTest, TestIndexScore) { ASSERT_EQ(SpecialFields::IndexScore().Type()->id(), arrow::Type::FLOAT); } +TEST(SpecialFieldsTest, TestRealtimeOffset) { + ASSERT_EQ(SpecialFields::RealtimeOffset().Id(), SpecialFieldIds::REALTIME_OFFSET); + ASSERT_EQ(SpecialFields::RealtimeOffset().Name(), "_REALTIME_OFFSET"); + ASSERT_EQ(SpecialFields::RealtimeOffset().Type()->id(), arrow::Type::INT64); + ASSERT_FALSE(SpecialFields::RealtimeOffset().Nullable()); +} + TEST(SpecialFieldsTest, TestKeyValueSpecialFieldCount) { ASSERT_EQ(SpecialFields::KEY_VALUE_SPECIAL_FIELD_COUNT, 2); } @@ -66,6 +73,7 @@ TEST(SpecialFieldsTest, TestIsSystemField) { ASSERT_TRUE(SpecialFields::IsSystemField("rowkind")); ASSERT_TRUE(SpecialFields::IsSystemField("_ROW_ID")); ASSERT_TRUE(SpecialFields::IsSystemField("_INDEX_SCORE")); + ASSERT_TRUE(SpecialFields::IsSystemField("_REALTIME_OFFSET")); ASSERT_TRUE(SpecialFields::IsSystemField("_KEY_0")); } diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index c6a07b2bf..859a57718 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -355,6 +355,22 @@ void ArrowUtils::TraverseArray(const std::shared_ptr& array) { } } +uint64_t ArrowUtils::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; +} + bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type) { if (type->id() != other_type->id() || type->num_fields() != other_type->num_fields()) { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index 13bd81549..a31d66de8 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include "arrow/api.h" @@ -48,6 +49,8 @@ class PAIMON_EXPORT ArrowUtils { // avoid subsequent multi-threading problems. static void TraverseArray(const std::shared_ptr& array); + static uint64_t GetArrayMemoryUsage(const std::shared_ptr& data); + static Result> RemoveFieldFromStructArray( const std::shared_ptr& struct_array, const std::string& field_name); diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 7e8986be5..1e9695332 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -24,12 +24,31 @@ #include #include +#include "arrow/c/abi.h" +#include "arrow/c/helpers.h" #include "arrow/memory_pool.h" #include "arrow/status.h" #include "fmt/format.h" #include "paimon/memory/memory_pool.h" namespace paimon { +namespace { + +struct ArrowArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleaseArrowArray(ArrowArray* array) { + std::unique_ptr data( + static_cast(array->private_data)); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); +} + +} // namespace class ArrowMemPoolAdaptor : public arrow::MemoryPool { public: @@ -107,4 +126,26 @@ std::unique_ptr GetArrowPool(const std::shared_ptr(pool); } +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release) { + return Status::Invalid("cannot retain Arrow array memory pool"); + } + if (!arrow_pool) { + ArrowArrayRelease(array); + return Status::Invalid("cannot retain Arrow array memory pool"); + } + std::unique_ptr data; + try { + data = std::make_unique( + ArrowArrayPrivateData{array->release, array->private_data, arrow_pool}); + } catch (const std::bad_alloc&) { + ArrowArrayRelease(array); + return Status::OutOfMemory("failed to retain Arrow array memory pool"); + } + array->private_data = data.release(); + array->release = ReleaseArrowArray; + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 96b59e3e8..214bb4509 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -23,11 +23,17 @@ #include "arrow/memory_pool.h" #include "paimon/memory/memory_pool.h" +#include "paimon/status.h" #include "paimon/visibility.h" +struct ArrowArray; + namespace paimon { PAIMON_EXPORT std::unique_ptr GetArrowPool( const std::shared_ptr& pool); +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool); + } // namespace paimon diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 1b6b71c69..858e302e8 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -20,6 +20,7 @@ #include #include +#include #include "arrow/api.h" #include "arrow/array/array_nested.h" diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 3b6806c73..dde9aaeba 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -154,6 +154,63 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReadersToFiles( + std::vector>&& readers) { + auto raw_readers_guard = ScopeGuard([&]() -> void { + for (std::unique_ptr& reader : readers) { + if (reader != nullptr) { + reader->Close(); + } + } + }); + if (readers.empty()) { + return Status::Invalid("sorted readers must not be empty"); + } + for (const std::unique_ptr& reader : readers) { + if (reader == nullptr) { + return Status::Invalid("sorted readers must not contain null reader"); + } + } + + // prepare loser tree sort merge reader + auto sort_merge_reader = std::make_unique( + std::move(readers), key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_); + raw_readers_guard.Release(); + // project key value to arrow array + auto create_consumer = [target_schema = write_schema_, pool = pool_]() + -> Result>> { + return KeyValueMetaProjectionConsumer::Create(target_schema, pool); + }; + // consumer batch size is WriteBatchSize + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), + /*projection_thread_num=*/1, pool_); + ScopeGuard async_readers_guard([&]() -> void { async_key_value_producer_consumer->Close(); }); + std::unique_ptr>> rolling_writer; + PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); + ScopeGuard abort_writer_guard([&]() -> void { rolling_writer->Abort(); }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_key_value_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + } + PAIMON_RETURN_NOT_OK(rolling_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, + rolling_writer->GetResult()); + abort_writer_guard.Release(); + + for (const std::shared_ptr& flushed_file : flushed_files) { + new_files_.emplace_back(flushed_file); + PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); + } + metrics_->Merge(rolling_writer->GetMetrics()); + return Status::OK(); +} + Status MergeTreeWriter::Compact(bool full_compaction) { return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction); } @@ -259,46 +316,7 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, // 1. flush write buffer to get sorted readers PAIMON_ASSIGN_OR_RAISE(std::vector> readers, write_buffer_->CreateReaders()); - // 2. prepare loser tree sort merge reader - auto sort_merge_reader = std::make_unique( - std::move(readers), key_comparator_, user_defined_seq_comparator_, - merge_function_wrapper_); - // 3. project key value to arrow array - auto create_consumer = [target_schema = write_schema_, pool = pool_]() - -> Result>> { - return KeyValueMetaProjectionConsumer::Create(target_schema, pool); - }; - // consumer batch size is WriteBatchSize - auto async_key_value_producer_consumer = - std::make_unique>( - std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), - /*projection_thread_num=*/1, pool_); - std::unique_ptr>> - rolling_writer; - PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); - ScopeGuard write_guard([&]() -> void { - rolling_writer->Abort(); - async_key_value_producer_consumer->Close(); - }); - while (true) { - PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, - async_key_value_producer_consumer->NextBatch()); - if (key_value_batch.batch == nullptr) { - break; - } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); - } - PAIMON_RETURN_NOT_OK(rolling_writer->Close()); - PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, - rolling_writer->GetResult()); - async_key_value_producer_consumer->Close(); - write_guard.Release(); - - for (const auto& flushed_file : flushed_files) { - new_files_.emplace_back(flushed_file); - PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); - } - metrics_->Merge(rolling_writer->GetMetrics()); + PAIMON_RETURN_NOT_OK(WriteSortedReadersToFiles(std::move(readers))); } PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction)); PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction)); diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index febce2afb..c2a9131e9 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.h +++ b/src/paimon/core/mergetree/merge_tree_writer.h @@ -51,6 +51,7 @@ class IOManager; class FieldsComparator; class MemoryPool; class Metrics; +class KeyValueRecordReader; template class MergeFunctionWrapper; @@ -69,6 +70,10 @@ class MergeTreeWriter : public BatchWriter { Status Write(std::unique_ptr&& batch) override; + /// Consumes readers whose complete streams are individually sorted by primary key and sequence + /// number. Readers are closed on success or failure. + Status WriteSortedReadersToFiles(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index c5a114f36..db2533fae 100644 --- a/src/paimon/core/mergetree/merge_tree_writer_test.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer_test.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -42,6 +43,7 @@ #include "paimon/core/io/compact_increment.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/data_increment.h" +#include "paimon/core/io/key_value_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" @@ -52,6 +54,8 @@ #include "paimon/fs/file_system.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/io_exception_helper.h" #include "paimon/testing/utils/read_result_collector.h" @@ -64,6 +68,36 @@ class MergeFunctionWrapper; } // namespace paimon namespace paimon::test { +namespace { + +class TrackingKeyValueRecordReader : public KeyValueRecordReader { + public: + TrackingKeyValueRecordReader(std::unique_ptr&& inner_reader, + bool* closed_flag) + : inner_reader_(std::move(inner_reader)), closed_flag_(closed_flag) {} + + Result> NextBatch() override { + return inner_reader_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return inner_reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_flag_ != nullptr) { + *closed_flag_ = true; + } + inner_reader_->Close(); + } + + private: + std::unique_ptr inner_reader_; + bool* closed_flag_; +}; + +} // namespace + class MergeTreeWriterTest : public ::testing::TestWithParam { public: class FakeCompactManager : public paimon::CompactManager { @@ -212,6 +246,23 @@ class MergeTreeWriterTest : public ::testing::TestWithParam { writer_compact_manager, io_manager, /*enable_multi_thread_spill=*/false, pool_); } + std::unique_ptr CreateSingleReader( + const std::shared_ptr& array, int32_t batch_size = 16, + const Status& next_batch_status = Status::OK()) const { + std::vector write_fields = {SpecialFields::SequenceNumber(), + SpecialFields::ValueKind()}; + write_fields.insert(write_fields.end(), value_fields_.begin(), value_fields_.end()); + std::shared_ptr write_schema = + DataField::ConvertDataFieldsToArrowSchema(write_fields); + std::shared_ptr key_schema = + arrow::schema(arrow::FieldVector({write_schema->field(2)})); + auto file_batch_reader = + std::make_unique(array, array->type(), batch_size); + file_batch_reader->SetNextBatchStatus(next_batch_status); + return std::make_unique( + std::move(file_batch_reader), key_schema, value_schema_, 0, pool_); + } + private: std::shared_ptr pool_; std::shared_ptr file_system_; @@ -377,6 +428,171 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); } +TEST_P(MergeTreeWriterTest, TestSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + std::string uuid = path_factory->uuid_; + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(4, dir->Str(), path_factory, 7, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])") + .ValueOrDie()); + + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + + ASSERT_OK(merge_writer->WriteSortedReadersToFiles(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + std::string expected_data_file_name = "data-" + uuid + "-0.orc"; + std::string expected_data_file_path = dir->Str() + "/" + expected_data_file_name; + ASSERT_OK_AND_ASSIGN(FileStatus data_file_status, + options.GetFileSystem()->GetFileStatus(expected_data_file_path)); + + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])"}, + &expected_array) + .ok()); + CheckFileContent(expected_data_file_path, expected_array); + + ASSERT_TRUE(commit_increment.GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(expected_data_file_name, new_file->file_name); + ASSERT_EQ(data_file_status.GetLen(), new_file->file_size); + ASSERT_EQ(3, new_file->row_count); + ASSERT_EQ(7, new_file->min_sequence_number); + ASSERT_EQ(9, new_file->max_sequence_number); + ASSERT_EQ(7, new_file->schema_id); + ASSERT_EQ(1, new_file->delete_row_count); +} + +TEST_P(MergeTreeWriterTest, TestMergeSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/4, dir->Str(), path_factory, + /*schema_id=*/7, options)); + + auto first_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [5, 0, "Alice", 10, 0, 15.1], + [7, 0, "Carol", 20, 1, 17.1], + [10, 0, "Eve", 30, 2, 20.1] + ])") + .ValueOrDie()); + auto second_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1] + ])") + .ValueOrDie()); + bool first_closed = false; + bool second_closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(first_array), &first_closed)); + sorted_readers.push_back(std::make_unique( + CreateSingleReader(second_array), &second_closed)); + + ASSERT_OK(merge_writer->WriteSortedReadersToFiles(std::move(sorted_readers))); + ASSERT_TRUE(first_closed); + ASSERT_TRUE(second_closed); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(5, new_file->row_count); + ASSERT_EQ(1, new_file->delete_row_count); + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [5, 0, "Alice", 10, 0, 15.1], + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1], + [10, 0, "Eve", 30, 2, 20.1] + ])"}, + &expected_array) + .ok()); + CheckFileContent(path_factory->ToPath(new_file), expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + std::vector> empty_readers; + Status empty_status = merge_writer->WriteSortedReadersToFiles(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReadersToFiles(std::move(null_readers)); + ASSERT_TRUE(null_status.IsInvalid()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + bool first_mixed_reader_closed = false; + bool second_mixed_reader_closed = false; + std::vector> mixed_readers; + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &first_mixed_reader_closed)); + mixed_readers.push_back(nullptr); + mixed_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &second_mixed_reader_closed)); + Status mixed_status = merge_writer->WriteSortedReadersToFiles(std::move(mixed_readers)); + ASSERT_TRUE(mixed_status.IsInvalid()); + ASSERT_TRUE(first_mixed_reader_closed); + ASSERT_TRUE(second_mixed_reader_closed); + + Status expected_status = Status::IOError("sorted reader failure"); + bool failing_reader_closed = false; + std::vector> failing_readers; + failing_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array, /*batch_size=*/16, expected_status), + &failing_reader_closed)); + Status failing_status = merge_writer->WriteSortedReadersToFiles(std::move(failing_readers)); + ASSERT_EQ(expected_status, failing_status); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); +} + TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({ diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 6807ae35e..fa1294d1b 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -57,6 +57,27 @@ struct KeyValue; template class MergeFunctionWrapper; +namespace { + +Status RestoreRealtimeCommittedProgress(const std::shared_ptr& realtime_context, + const std::shared_ptr& snapshot_manager, + const CoreOptions& options) { + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager->LatestSnapshot()); + if (latest_snapshot) { + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap realtime_committed_offsets, + RealtimeCommitProperties::ReadOffsets(latest_snapshot, options.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( + latest_snapshot->Id(), realtime_committed_offsets)); + } + return Status::OK(); +} + +} // namespace + Result> FileStoreWrite::PrepareCommitWithProgress(int64_t) { return Status::Invalid("prepare commit with progress requires a real-time writer"); } @@ -143,17 +164,8 @@ Result> FileStoreWrite::Create(std::unique_ptr latest_snapshot, - snapshot_manager->LatestSnapshot()); - if (latest_snapshot) { - PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap realtime_committed_offsets, - RealtimeCommitProperties::ReadOffsets( - latest_snapshot, options.GetFileSystem())); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(ctx->GetRealtimeContext())); - PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( - latest_snapshot->Id(), realtime_committed_offsets)); - } + PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(), + snapshot_manager, options)); } std::shared_ptr write_schema = arrow_schema; const auto& write_field_names = ctx->GetWriteSchema(); @@ -197,7 +209,16 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - return Status::Invalid("real-time write currently supports append tables only"); + PAIMON_RETURN_NOT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *schema)); + if (ignore_previous_files) { + return Status::NotImplemented( + "PK realtime requires restore from the latest snapshot"); + } + if (!ctx->GetWriteSchema().empty()) { + return Status::NotImplemented("PK realtime does not support a custom write schema"); + } + PAIMON_RETURN_NOT_OK(RestoreRealtimeCommittedProgress(ctx->GetRealtimeContext(), + snapshot_manager, options)); } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( @@ -253,7 +274,8 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory, io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options, ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), - ctx->EnableMultiThreadSpill(), ctx->GetExecutor(), ctx->GetMemoryPool()); + ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), + ctx->GetMemoryPool()); } } diff --git a/src/paimon/core/operation/key_value_file_store_write.cpp b/src/paimon/core/operation/key_value_file_store_write.cpp index 08c5ea0c3..152a4ed01 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -18,21 +18,30 @@ #include "paimon/core/operation/key_value_file_store_write.h" +#include #include +#include "arrow/c/bridge.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_list.h" #include "paimon/core/mergetree/levels.h" #include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/realtime/realtime_primary_key_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" +#include "paimon/realtime/realtime_context.h" namespace arrow { class Schema; @@ -60,6 +69,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, bool enable_multi_thread_spill, + const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool) : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user, root_path, table_schema, schema, /*write_schema=*/schema, @@ -67,6 +77,7 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( ignore_previous_files, is_streaming_mode, ignore_num_bucket_check, executor, pool), enable_multi_thread_spill_(enable_multi_thread_spill), + realtime_context_(realtime_context), key_comparator_(key_comparator), user_defined_seq_comparator_(user_defined_seq_comparator), merge_function_wrapper_(merge_function_wrapper), @@ -74,7 +85,11 @@ KeyValueFileStoreWrite::KeyValueFileStoreWrite( options_, key_comparator_, user_defined_seq_comparator_, compaction_metrics_, table_schema_, schema_, schema_manager_, io_manager_, cache_manager_, file_store_path_factory_, root_path_, pool_)), - logger_(Logger::GetLogger("KeyValueFileStoreWrite")) {} + logger_(Logger::GetLogger("KeyValueFileStoreWrite")) { + if (realtime_context_) { + writer_memory_manager_ = std::make_unique(); + } +} Result> KeyValueFileStoreWrite::CreateFileStoreScan( const std::shared_ptr& scan_filter) const { @@ -106,22 +121,68 @@ Result> KeyValueFileStoreWrite::CreateWriter( file_store_path_factory_->CreateDataFilePathFactory(partition, bucket)); PAIMON_ASSIGN_OR_RAISE(std::vector trimmed_primary_keys, table_schema_->TrimmedPrimaryKeys()); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr levels, - Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); - auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr compact_manager, - compact_manager_factory_->CreateCompactManager(partition, bucket, compact_strategy, - compact_executor_, levels, dv_maintainer)); + std::map partition_map; + std::shared_ptr compact_manager; + std::shared_ptr realtime_context_impl; + std::optional realtime_store_state; + std::shared_ptr transport_schema; + if (realtime_context_) { + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + file_store_path_factory_->GeneratePartitionVector(partition)); + partition_map = + std::map(partition_values.begin(), partition_values.end()); + PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); + transport_schema = RealtimePrimaryKeyLayout::CreateSchema(schema_->fields()); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*transport_schema, c_write_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(c_write_schema), options_.ToMap(), pool_, + RealtimeStoreMode::PRIMARY_KEY}, + RealtimePartitionBucket(partition_map, bucket))); + realtime_store_state = std::move(store_state); + compact_manager = std::make_shared(); + } else { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr levels, + Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); + auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); + PAIMON_ASSIGN_OR_RAISE(compact_manager, compact_manager_factory_->CreateCompactManager( + partition, bucket, compact_strategy, + compact_executor_, levels, dv_maintainer)); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, - options_, compact_manager, io_manager_, enable_multi_thread_spill_, pool_)); - return writer; + options_, compact_manager, realtime_context_ ? nullptr : io_manager_, + enable_multi_thread_spill_, pool_)); + if (!realtime_context_) { + return std::shared_ptr(std::move(writer)); + } + return RealtimePrimaryKeyWriter::Create(partition_map, bucket, schema_, transport_schema, + trimmed_primary_keys, key_comparator_, options_, + realtime_context_impl, realtime_store_state.value(), + restore_max_seq_number, writer, pool_); +} + +Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + if (!realtime_context_) { + return Status::Invalid("refresh committed snapshot requires a real-time writer"); + } + PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, snapshot_manager_->LoadSnapshot(snapshot_id)); + PAIMON_ASSIGN_OR_RAISE( + RealtimeOffsetMap committed_offsets, + RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), + options_.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context_)); + return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets); } Status KeyValueFileStoreWrite::Close() { diff --git a/src/paimon/core/operation/key_value_file_store_write.h b/src/paimon/core/operation/key_value_file_store_write.h index 14457590f..66c362f2e 100644 --- a/src/paimon/core/operation/key_value_file_store_write.h +++ b/src/paimon/core/operation/key_value_file_store_write.h @@ -45,6 +45,7 @@ class SnapshotManager; class SchemaManager; class TableSchema; class IOManager; +class RealtimeContext; struct KeyValue; template class MergeFunctionWrapper; @@ -65,8 +66,10 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { const std::shared_ptr>& merge_function_wrapper, const CoreOptions& options, bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check, bool enable_multi_thread_spill, + const std::shared_ptr& realtime_context, const std::shared_ptr& executor, const std::shared_ptr& pool); + Status RefreshCommittedSnapshot(int64_t snapshot_id) override; Status Close() override; private: @@ -79,8 +82,13 @@ class KeyValueFileStoreWrite : public AbstractFileStoreWrite { Result> CreateFileStoreScan( const std::shared_ptr& filter) const override; + bool IsRealtimeWrite() const override { + return realtime_context_ != nullptr; + } + private: bool enable_multi_thread_spill_; + std::shared_ptr realtime_context_; std::shared_ptr key_comparator_; std::shared_ptr user_defined_seq_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index 35d938af7..88b848eea 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -19,9 +19,13 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include +#include #include #include +#include +#include #include +#include #include #include @@ -40,10 +44,13 @@ #include "paimon/common/data/shredding/map_shared_shredding_utils.h" #include "paimon/common/data/shredding/map_shredding_defs.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/operation/restore_files.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/file_store_commit.h" @@ -52,7 +59,9 @@ #include "paimon/format/file_format_factory.h" #include "paimon/format/reader_builder.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" #include "paimon/testing/utils/test_helper.h" @@ -60,6 +69,50 @@ #include "paimon/write_context.h" namespace paimon::test { +namespace { + +class TestingMemoryPool final : public MemoryPool { + public: + void* Malloc(uint64_t size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* pointer, size_t old_size, size_t new_size, uint64_t alignment) override { + ++allocation_count; + if (reject_allocations) { + throw std::bad_alloc(); + } + return delegate_->Realloc(pointer, old_size, new_size, alignment); + } + + void Free(void* pointer, uint64_t size) override { + delegate_->Free(pointer, size); + } + + void Free(void* pointer, uint64_t size, uint64_t alignment) override { + delegate_->Free(pointer, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + bool reject_allocations = false; + int64_t allocation_count = 0; + + private: + std::unique_ptr delegate_ = GetMemoryPool(); +}; + +} // namespace class KeyValueFileStoreWriteTest : public ::testing::Test { protected: @@ -127,14 +180,15 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } - std::unique_ptr MakeBatch(const std::shared_ptr& schema, - const std::string& json) const { + std::unique_ptr MakeBatch( + const std::shared_ptr& schema, const std::string& json, + const std::vector& row_kinds = {}) const { auto struct_type = arrow::struct_(schema->fields()); auto array = arrow::ipc::internal::json::ArrayFromJSON(struct_type, json).ValueOrDie(); ::ArrowArray arrow_array; EXPECT_TRUE(arrow::ExportArray(*array, &arrow_array).ok()); RecordBatchBuilder batch_builder(&arrow_array); - return batch_builder.SetBucket(0).Finish().value(); + return batch_builder.SetRowKinds(row_kinds).SetBucket(0).Finish().value(); } std::vector> WriteAndPrepare( @@ -193,6 +247,67 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { EXPECT_NE(nullptr, metadata); return MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy()).value(); } + + Result>> + ReadRealtimePrimaryKeyTransportRows( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + context->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected exactly one real-time store"); + } + arrow::FieldVector value_fields = {DataField::ConvertDataFieldToArrowField(DataField( + 0, arrow::field("id", arrow::int64(), false))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}; + std::shared_ptr transport_schema = + RealtimePrimaryKeyLayout::CreateSchema(value_fields); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*transport_schema, c_schema.get())); + RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + views[0].store->CreateQueryReaders(views[0].read_view, 0, query_context)); + std::vector> rows; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(array); + if (!values || values->num_fields() != 5) { + return Status::Invalid("unexpected realtime primary-key transport batch"); + } + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(values->field(0)); + std::shared_ptr sequences = + std::dynamic_pointer_cast(values->field(1)); + std::shared_ptr offsets = + std::dynamic_pointer_cast(values->field(2)); + std::shared_ptr ids = + std::dynamic_pointer_cast(values->field(3)); + std::shared_ptr payloads = + std::dynamic_pointer_cast(values->field(4)); + if (!row_kinds || !sequences || !offsets || !ids || !payloads) { + return Status::Invalid("unexpected realtime primary-key transport column type"); + } + for (int64_t row = 0; row < values->length(); ++row) { + rows.emplace_back(row_kinds->Value(row), ids->Value(row), + payloads->GetString(row), sequences->Value(row), + offsets->Value(row)); + } + } + reader->Close(); + } + return rows; + } }; TEST_F(KeyValueFileStoreWriteTest, TestWriteWithInvalidBatch) { @@ -303,6 +418,197 @@ TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenLookupEnabl ASSERT_EQ(commit_messages.size(), 1); } +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeWrite) { + const std::map options = { + {Options::BUCKET, "1"}, + {Options::WRITE_BUFFER_SIZE, "1"}, + {Options::REALTIME_ENABLED, "true"}, + }; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithTempDirectory(dir->Str()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + std::unique_ptr batch = + MakeBatch(schema, R"([ + [1, "old"], + [2, "two"], + [1, "new"] + ])", + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER}); + ASSERT_OK(writer->Write(std::move(batch))); + using RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector transport_rows, + ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{ + {0, 1, "old", 0, 0}, {2, 1, "new", 2, 2}, {3, 2, "two", 1, 1}}), + transport_rows); + ASSERT_OK_AND_ASSIGN(std::vector progresses, + writer->PrepareCommitWithProgress(0)); + ASSERT_EQ(1, progresses.size()); + ASSERT_EQ(OffsetRange(0, 3), progresses[0].offset_range); + std::shared_ptr commit_message = + std::dynamic_pointer_cast(progresses[0].commit_message); + ASSERT_NE(nullptr, commit_message); + int64_t row_count = 0; + for (const std::shared_ptr& file : + commit_message->GetNewFilesIncrement().NewFiles()) { + row_count += file->row_count; + } + ASSERT_EQ(2, row_count); + ASSERT_EQ(0, TestHelper::CountChannelFiles(dir->GetFileSystem(), dir->Str())); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimePool) { + const std::map options = {{Options::BUCKET, "1"}, + {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + std::shared_ptr pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "test"); + builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + const int64_t allocations_before_write = pool->allocation_count; + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "one"]])"))); + ASSERT_GT(pool->allocation_count, allocations_before_write); + ASSERT_OK(writer->Close()); + writer.reset(); + using RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector retained_rows, + ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "one", 0, 0}}), retained_rows); + + std::shared_ptr rejecting_pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rejecting_context, + RealtimeContext::Create()); + WriteContextBuilder rejecting_builder(table_path, "rejecting"); + rejecting_builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(rejecting_context) + .WithMemoryPool(rejecting_pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_write_context, + rejecting_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_writer, + FileStoreWrite::Create(std::move(rejecting_write_context))); + ASSERT_OK(rejecting_writer->Write(MakeBatch(schema, "[]"))); + const int64_t rejecting_allocations_before_write = rejecting_pool->allocation_count; + rejecting_pool->reject_allocations = true; + ASSERT_NOK_WITH_MSG(rejecting_writer->Write(MakeBatch(schema, R"([[2, "two"]])")), + "Out of memory"); + ASSERT_GT(rejecting_pool->allocation_count, rejecting_allocations_before_write); + ASSERT_OK_AND_ASSIGN(std::vector rejected_rows, + ReadRealtimePrimaryKeyTransportRows(rejecting_context)); + ASSERT_TRUE(rejected_rows.empty()); + ASSERT_OK(rejecting_writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { + const int64_t max = std::numeric_limits::max(); + const std::map options = {{Options::BUCKET, "1"}, + {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("value", arrow::utf8()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + const std::string table_path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr initial_context, + RealtimeContext::Create()); + WriteContextBuilder initial_builder(table_path, "initial"); + initial_builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext( + initial_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_write_context, + initial_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr initial_writer, + FileStoreWrite::Create(std::move(initial_write_context))); + ASSERT_OK(initial_writer->Write(MakeBatch(schema, R"([[0, "initial"]])"))); + ASSERT_OK_AND_ASSIGN(std::vector initial_progress, + initial_writer->PrepareCommitWithProgress(0)); + ASSERT_EQ(1, initial_progress.size()); + std::shared_ptr initial_message = + std::dynamic_pointer_cast(initial_progress[0].commit_message); + ASSERT_NE(nullptr, initial_message); + ASSERT_EQ(1, initial_message->GetNewFilesIncrement().NewFiles().size()); + initial_message->GetNewFilesIncrement().NewFiles()[0]->AssignSequenceNumber(max - 2, max - 2); + initial_progress[0].offset_range = OffsetRange(0, max - 1); + + CommitContextBuilder commit_builder(table_path, "initial"); + commit_builder.SetOptions(options); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, commit_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committer, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, + committer->CommitWithProgress(initial_progress, 0, std::nullopt)); + ASSERT_OK(initial_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(table_path, "boundary"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + ASSERT_OK(writer->Write(MakeBatch(schema, R"([[1, "legal"]])"))); + using RealtimePrimaryKeyTransportRow = + std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector transport_rows, + ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), + transport_rows); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[2, "overflow"]])")), + "real-time offset range exceeds INT64_MAX"); + ASSERT_OK_AND_ASSIGN(transport_rows, ReadRealtimePrimaryKeyTransportRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), + transport_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context_impl, + RealtimeContextImpl::Cast(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::vector views, + context_impl->AcquireReadViews()); + ASSERT_EQ(1, views.size()); + ASSERT_EQ(std::optional(OffsetRange(max - 1, max)), + views[0].read_view->GetOffsetRange()); + ASSERT_OK(writer->Close()); + ASSERT_GE(snapshot_id, 1); +} + TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenDefaultCompactRewriterPathEnabled) { ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index b753ea431..1f96c195c 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -36,6 +36,7 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/object_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/apply_deletion_vector_batch_reader.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" @@ -78,6 +79,185 @@ struct KeyValue; template class MergeFunctionWrapper; +namespace { + +class SortMergeKeyValueRecordReader : public KeyValueRecordReader { + public: + explicit SortMergeKeyValueRecordReader(std::unique_ptr&& reader) + : reader_(std::move(reader)) {} + + class Iterator : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(std::unique_ptr&& iterator) + : iterator_(std::move(iterator)) {} + + Result HasNext() const override { + return iterator_->HasNext(); + } + + Result Next() override { + return std::move(iterator_->Next()); + } + + private: + std::unique_ptr iterator_; + }; + + Result> NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iterator, + reader_->NextBatch()); + if (!iterator) { + return std::unique_ptr(); + } + return std::make_unique(std::move(iterator)); + } + + void Close() override { + reader_->Close(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + private: + std::unique_ptr reader_; +}; + +} // namespace + +class MergeFileSplitRead::RealtimeReaderBuilder { + public: + static Result> Create( + const std::vector>& disk_splits, + std::vector>&& additional_readers, + MergeFileSplitRead* owner) { + ScopeGuard additional_readers_guard([&additional_readers]() { + for (const std::unique_ptr& reader : additional_readers) { + if (reader) { + reader->Close(); + } + } + }); + RealtimeReaderBuilder builder(owner); + std::vector> readers; + if (!disk_splits.empty()) { + PAIMON_RETURN_NOT_OK(builder.CollectDiskReaders(disk_splits, &readers)); + } + readers.reserve(readers.size() + additional_readers.size()); + for (std::unique_ptr& additional_reader : additional_readers) { + readers.push_back(std::move(additional_reader)); + } + additional_readers_guard.Release(); + return builder.CreateMergedReader(std::move(readers)); + } + + private: + explicit RealtimeReaderBuilder(MergeFileSplitRead* owner) : owner_(owner) {} + + Status CollectDiskReaders(const std::vector>& disk_splits, + std::vector>* readers) { + std::shared_ptr first_split; + std::vector> data_files; + std::vector> deletion_files; + for (const std::shared_ptr& disk_split : disk_splits) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + if (!data_split) { + return Status::Invalid("merge input disk split is not a data split"); + } + if (!first_split) { + first_split = data_split; + } + const std::vector>& split_files = data_split->DataFiles(); + const std::vector>& split_deletion_files = + data_split->DeletionFiles(); + if (!split_deletion_files.empty() && + split_deletion_files.size() != split_files.size()) { + return Status::Invalid( + "merge input disk split deletion files must be empty or match data files"); + } + data_files.insert(data_files.end(), split_files.begin(), split_files.end()); + if (split_deletion_files.empty()) { + deletion_files.insert(deletion_files.end(), split_files.size(), std::nullopt); + } else { + deletion_files.insert(deletion_files.end(), split_deletion_files.begin(), + split_deletion_files.end()); + } + } + const BinaryRow& partition = first_split->Partition(); + const int32_t bucket = first_split->Bucket(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); + + DeletionVector::Factory dv_factory; + std::vector> disk_sections; + PAIMON_RETURN_NOT_OK( + owner_->CreateDiskSections(data_files, deletion_files, &dv_factory, &disk_sections)); + if (disk_sections.empty()) { + return Status::OK(); + } + std::vector> section_readers; + ScopeGuard section_readers_guard([§ion_readers]() { + for (const std::unique_ptr& reader : section_readers) { + reader->Close(); + } + }); + section_readers.reserve(disk_sections.size()); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr> merge_function_wrapper, + MergeFileSplitRead::CreateMergeFunctionWrapper(owner_->options_, + owner_->context_->GetTableSchema(), + owner_->value_schema_, owner_->pool_)); + for (const std::vector& section : disk_sections) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr section_reader, + owner_->CreateSortMergeReaderForSection( + section, partition, dv_factory, owner_->predicate_for_keys_, + data_file_path_factory, /*drop_delete=*/false, merge_function_wrapper)); + section_readers.push_back( + std::make_unique(std::move(section_reader))); + } + std::unique_ptr concat_reader = + std::make_unique(std::move(section_readers)); + section_readers_guard.Release(); + readers->push_back(std::move(concat_reader)); + return Status::OK(); + } + + Result> CreateMergedReader( + std::vector>&& record_readers) { + ScopeGuard record_readers_guard([&record_readers]() { + for (const std::unique_ptr& reader : record_readers) { + if (reader) { + reader->Close(); + } + } + }); + if (record_readers.empty()) { + record_readers_guard.Release(); + return std::make_unique(std::vector>{}, + owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + owner_->CreateSortMergeReader(std::move(record_readers))); + record_readers_guard.Release(); + ScopeGuard sort_merge_reader_guard([&sort_merge_reader]() { + if (sort_merge_reader) { + sort_merge_reader->Close(); + } + }); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr result, + owner_->CreateProjectedReader(std::move(sort_merge_reader), + owner_->context_->GetPredicate(), + /*complete_row_kind=*/true)); + sort_merge_reader_guard.Release(); + return result; + } + + MergeFileSplitRead* owner_; +}; + Result> MergeFileSplitRead::Create( const std::shared_ptr& path_factory, const std::shared_ptr& context, @@ -158,6 +338,12 @@ Result> MergeFileSplitRead::CreateReader( return std::make_unique(std::move(batch_reader), pool_); } +Result> MergeFileSplitRead::CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector>&& additional_readers) { + return RealtimeReaderBuilder::Create(disk_splits, std::move(additional_readers), this); +} + void MergeFileSplitRead::SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper) { merge_function_wrapper_ = merge_function_wrapper; @@ -236,13 +422,10 @@ Result> MergeFileSplitRead::ApplyIndexAndDvRead Result> MergeFileSplitRead::CreateMergeReader( const std::shared_ptr& data_split, const std::shared_ptr& data_file_path_factory) { - auto dv_factory = DeletionVector::CreateFactory( - options_.GetFileSystem(), - DeletionVector::CreateDeletionFileMap(data_split->DataFiles(), data_split->DeletionFiles()), - pool_); - - std::vector> sections = - IntervalPartition(data_split->DataFiles(), key_comparator_).Partition(); + DeletionVector::Factory dv_factory; + std::vector> sections; + PAIMON_RETURN_NOT_OK(CreateDiskSections(data_split->DataFiles(), data_split->DeletionFiles(), + &dv_factory, §ions)); std::vector> batch_readers; batch_readers.reserve(sections.size()); // no overlap through multiple sections @@ -453,38 +636,100 @@ Result> MergeFileSplitRead::CreateReaderForSection( } else { predicate = context_->GetPredicate(); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, - CreateSortMergeReaderForSection(section, partition, dv_factory, - predicate, data_file_path_factory, - /*drop_delete=*/!force_keep_delete_)); - // KeyValueProjectionReader converts KeyValue objects to arrow array according to projection - if (!context_->EnableMultiThreadRowToBatch()) { - return KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, - projection_, options_.GetReadBatchSize(), pool_); - } - int32_t thread_number = context_->GetRowToBatchThreadNumber(); - assert(thread_number > 0); - return std::make_unique( - std::move(sort_merge_reader), raw_read_schema_, projection_, options_.GetReadBatchSize(), - thread_number, pool_); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr sort_merge_reader, + CreateSortMergeReaderForSection(section, partition, dv_factory, predicate, + data_file_path_factory, /*drop_delete=*/false)); + return CreateProjectedReader(std::move(sort_merge_reader), /*predicate=*/nullptr, + /*complete_row_kind=*/false); } -Result> MergeFileSplitRead::CreateSortMergeReaderForSection( +Status MergeFileSplitRead::CreateDiskSections( + const std::vector>& data_files, + const std::vector>& deletion_files, + DeletionVector::Factory* dv_factory, std::vector>* sections) const { + *dv_factory = DeletionVector::CreateFactory( + options_.GetFileSystem(), DeletionVector::CreateDeletionFileMap(data_files, deletion_files), + pool_); + *sections = IntervalPartition(data_files, key_comparator_).Partition(); + return Status::OK(); +} + +Result>> +MergeFileSplitRead::CreateRecordReadersForSection( const std::vector& section, const BinaryRow& partition, DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, - const std::shared_ptr& data_file_path_factory, bool drop_delete) { - // with overlap in one section + const std::shared_ptr& data_file_path_factory) const { std::vector> record_readers; record_readers.reserve(section.size()); - for (const auto& run : section) { + for (const SortedRun& run : section) { // no overlap in a run PAIMON_ASSIGN_OR_RAISE( std::unique_ptr run_reader, CreateReaderForRun(partition, run, dv_factory, predicate, data_file_path_factory)); record_readers.emplace_back(std::move(run_reader)); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, - CreateSortMergeReader(std::move(record_readers))); + return record_readers; +} + +Result> MergeFileSplitRead::CreateProjectedReader( + std::unique_ptr&& sort_merge_reader, + const std::shared_ptr& predicate, bool complete_row_kind) { + if (!force_keep_delete_) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } + // KeyValueProjectionReader converts KeyValue objects to arrow array according to projection + std::unique_ptr projection_reader; + if (!context_->EnableMultiThreadRowToBatch()) { + PAIMON_ASSIGN_OR_RAISE( + projection_reader, + KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, + projection_, options_.GetReadBatchSize(), pool_)); + } else { + const int32_t thread_number = context_->GetRowToBatchThreadNumber(); + assert(thread_number > 0); + projection_reader = std::make_unique( + std::move(sort_merge_reader), raw_read_schema_, projection_, + options_.GetReadBatchSize(), thread_number, pool_); + } + ScopeGuard projection_reader_guard([&projection_reader]() { + if (projection_reader) { + projection_reader->Close(); + } + }); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr filtered_reader, + ApplyPredicateFilterIfNeeded(std::move(projection_reader), predicate)); + projection_reader_guard.Release(); + projection_reader = std::move(filtered_reader); + if (complete_row_kind) { + return std::make_unique(std::move(projection_reader), pool_); + } + return projection_reader; +} + +Result> MergeFileSplitRead::CreateSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, bool drop_delete) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr> merge_function_wrapper, + GetMergeFunctionWrapper()); + return CreateSortMergeReaderForSection(section, partition, dv_factory, predicate, + data_file_path_factory, drop_delete, + merge_function_wrapper); +} + +Result> MergeFileSplitRead::CreateSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, bool drop_delete, + const std::shared_ptr>& merge_function_wrapper) { + // with overlap in one section + PAIMON_ASSIGN_OR_RAISE(std::vector> record_readers, + CreateRecordReadersForSection(section, partition, dv_factory, predicate, + data_file_path_factory)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr sort_merge_reader, + CreateSortMergeReader(std::move(record_readers), merge_function_wrapper)); if (drop_delete) { sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); } @@ -519,6 +764,12 @@ Result> MergeFileSplitRead::CreateSortMergeRead std::vector>&& record_readers) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr> merge_function_wrapper, GetMergeFunctionWrapper()); + return CreateSortMergeReader(std::move(record_readers), merge_function_wrapper); +} + +Result> MergeFileSplitRead::CreateSortMergeReader( + std::vector>&& record_readers, + const std::shared_ptr>& merge_function_wrapper) const { auto sort_engine = options_.GetSortEngine(); if (sort_engine == SortEngine::MIN_HEAP) { return std::make_unique( diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 11dcd0b37..36ba35dda 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -117,10 +117,20 @@ class MergeFileSplitRead : public AbstractSplitRead { return value_schema_; } + std::shared_ptr GetKeySchema() const { + return key_schema_; + } + + Result> CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector>&& additional_readers); + void SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper); private: + class RealtimeReaderBuilder; + Result> CreateMergeReader( const std::shared_ptr& data_split, const std::shared_ptr& data_file_path_factory); @@ -134,6 +144,26 @@ class MergeFileSplitRead : public AbstractSplitRead { DeletionVector::Factory dv_factory, const std::shared_ptr& data_file_path_factory); + Status CreateDiskSections(const std::vector>& data_files, + const std::vector>& deletion_files, + DeletionVector::Factory* dv_factory, + std::vector>* sections) const; + + Result>> CreateRecordReadersForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory) const; + + Result> CreateProjectedReader( + std::unique_ptr&& sort_merge_reader, + const std::shared_ptr& predicate, bool complete_row_kind); + + Result> CreateSortMergeReaderForSection( + const std::vector& section, const BinaryRow& partition, + DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, + const std::shared_ptr& data_file_path_factory, bool drop_delete, + const std::shared_ptr>& merge_function_wrapper); + Result> CreateReaderForRun( const BinaryRow& partition, const SortedRun& sorted_run, DeletionVector::Factory dv_factory, const std::shared_ptr& predicate, @@ -142,6 +172,10 @@ class MergeFileSplitRead : public AbstractSplitRead { Result> CreateSortMergeReader( std::vector>&& record_readers); + Result> CreateSortMergeReader( + std::vector>&& record_readers, + const std::shared_ptr>& merge_function_wrapper) const; + Result>> GetMergeFunctionWrapper(); MergeFileSplitRead(const std::shared_ptr& path_factory, diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index f857d4876..d02120de4 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -40,6 +40,7 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/key_value_in_memory_record_reader.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/schema/schema_manager.h" @@ -51,7 +52,6 @@ #include "paimon/executor.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" -#include "paimon/metrics.h" #include "paimon/predicate/literal.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" @@ -66,6 +66,7 @@ class FileSystem; } // namespace paimon namespace paimon::test { + // Parameter: min_heap/loser_tree; enable/disable IO prefetch; enable/disable multi thread row to // batch class MergeFileSplitReadTest : public ::testing::Test, @@ -328,9 +329,8 @@ class MergeFileSplitReadTest : public ::testing::Test, return {data_split1}; } - Result> CreateReader( - const std::shared_ptr& internal_context, - const std::vector>& data_splits) { + Result> CreateMergeFileSplitRead( + const std::shared_ptr& internal_context) { const auto& core_options = internal_context->GetCoreOptions(); const auto& table_schema = internal_context->GetTableSchema(); auto arrow_schema = DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); @@ -347,9 +347,14 @@ class MergeFileSplitReadTest : public ::testing::Test, core_options.DataFilePrefix(), core_options.LegacyPartitionNameEnabled(), external_paths, global_index_external_path, core_options.IndexFileInDataFileDir(), pool_)); - PAIMON_ASSIGN_OR_RAISE(auto split_read, - MergeFileSplitRead::Create(path_factory, std::move(internal_context), - pool_, executor_)); + return MergeFileSplitRead::Create(path_factory, internal_context, pool_, executor_); + } + + Result> CreateReader( + const std::shared_ptr& internal_context, + const std::vector>& data_splits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr split_read, + CreateMergeFileSplitRead(internal_context)); std::vector> batch_readers; batch_readers.reserve(data_splits.size()); for (const auto& split : data_splits) { @@ -666,6 +671,73 @@ TEST_P(MergeFileSplitReadTest, TestSimple) { CheckResult(result_array, expected_array, read_schema); } +TEST_P(MergeFileSplitReadTest, TestRealtimeReadConcatenatesOrderedDiskSections) { + std::string path = + paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; + ReadContextBuilder context_builder(path); + std::vector raw_read_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("k1", arrow::int32())), + DataField(5, arrow::field("s1", arrow::utf8())), + DataField(6, arrow::field("v0", arrow::float64()))}; + std::shared_ptr read_schema = + DataField::ConvertDataFieldsToArrowSchema(raw_read_fields); + ASSERT_TRUE(read_schema); + + context_builder.SetReadFieldNames({"k0", "k1", "s1", "v0"}); + context_builder.SetOptions( + {{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::MERGE_ENGINE, "deduplicate"}}); + AddOptions(&context_builder); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); + std::shared_ptr internal_context = CreateInternalReadContext(read_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr split_read, + CreateMergeFileSplitRead(internal_context)); + + std::shared_ptr memory_type = + arrow::struct_(split_read->GetValueSchema()->fields()); + std::shared_ptr memory_array = + std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(memory_type, R"([ + [100, 200, "memory-late", 10000.0, "zzzz"], + [1, 1, "memory-delete", 1100.0, "zzzz"], + [0, 0, "memory-first", 1000.0, "zzzz"], + [50, 0, "memory-middle", 5000.0, "zzzz"] + ])") + .ValueOrDie()); + std::vector> memory_readers; + memory_readers.push_back(std::make_unique( + /*last_sequence_num=*/9, memory_array, + std::vector( + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT}), + std::vector({"k0", "k1"}), std::vector({"s0", "s1"}), + /*sequence_fields_ascending=*/true, split_read->GetKeyComparator(), pool_)); + + std::vector> disk_splits = {PrepareDataSplit().front()}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch_reader, + split_read->CreateRealtimeReader(disk_splits, std::move(memory_readers))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, + ReadResultCollector::CollectResult(batch_reader.get())); + + arrow::FieldVector fields_with_row_kind = read_schema->fields(); + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + std::shared_ptr expected_array; + auto expected_status = + arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(fields_with_row_kind), {R"([ + [0, 0, 0, "memory-first", 1000.0], + [0, 0, 1, "you", 11.1], + [0, 1, 0, "later", 12.2], + [0, 1, 2, "!", 13.3], + [0, 50, 0, "memory-middle", 5000.0], + [0, 100, 200, "memory-late", 10000.0] + ])"}, + &expected_array); + ASSERT_TRUE(expected_status.ok()); + CheckResult(result_array, expected_array, read_schema); + ASSERT_TRUE(batch_reader->GetReaderMetrics()); + batch_reader->Close(); +} + TEST_P(MergeFileSplitReadTest, TestLookUp) { std::string path = paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index 18087a29f..1136243e8 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -33,6 +33,7 @@ #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/projected_array.h" @@ -43,22 +44,6 @@ 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; -} - bool SupportsMinMax(const std::shared_ptr& type) { switch (type->id()) { case arrow::Type::BOOL: @@ -393,11 +378,11 @@ Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { if (building_range_ && write_batch.offset_range.begin != building_range_->end) { return Status::Invalid("real-time offset ranges must be contiguous"); } - uint64_t memory_usage = GetArrayMemoryUsage(struct_array->data()); + uint64_t memory_usage = ArrowUtils::GetArrayMemoryUsage(struct_array->data()); if (statistics) { - memory_usage += GetArrayMemoryUsage(statistics->min_values->data()) + - GetArrayMemoryUsage(statistics->max_values->data()) + - GetArrayMemoryUsage(statistics->null_counts->data()); + memory_usage += ArrowUtils::GetArrayMemoryUsage(statistics->min_values->data()) + + ArrowUtils::GetArrayMemoryUsage(statistics->max_values->data()) + + ArrowUtils::GetArrayMemoryUsage(statistics->null_counts->data()); } building_memory_usage_ += memory_usage; building_batches_.push_back( diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 1d7219c41..dff12b589 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -25,25 +25,37 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" namespace paimon { Result> ArrowRealtimeStoreFactory::Create( - std::unique_ptr write_schema, StatisticsMode statistics_mode, - const std::map&, const std::shared_ptr& memory_pool) { - if (!write_schema || !write_schema->release) { + RealtimeStoreCreateRequest&& request) { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("real-time store write schema is null"); } - ScopeGuard schema_guard([schema = write_schema.get()]() { ArrowSchemaRelease(schema); }); - if (!memory_pool) { + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + if (!request.memory_pool) { return Status::Invalid("real-time store memory pool is null"); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, - arrow::ImportSchema(write_schema.get())); - std::shared_ptr arrow_pool = GetArrowPool(memory_pool); - return std::make_shared(imported_schema, statistics_mode, memory_pool, - arrow_pool); + arrow::ImportSchema(request.write_schema.get())); + switch (request.mode) { + case RealtimeStoreMode::APPEND_ONLY: { + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + return std::make_shared(imported_schema, request.statistics_mode, + request.memory_pool, arrow_pool); + } + case RealtimeStoreMode::PRIMARY_KEY: { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool)); + return std::shared_ptr(std::move(store)); + } + } + return Status::Invalid("invalid real-time store mode: ", static_cast(request.mode)); } } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index 9aae99332..864b1f810 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -232,8 +232,11 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { ArrowRealtimeStoreFactory factory; std::unique_ptr write_schema = MakeReadSchema(schema_); + RealtimeStoreCreateRequest request{std::move(write_schema), + /*options=*/{}, pool_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::FULL}; ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_store, - factory.Create(std::move(write_schema), StatisticsMode::FULL, {}, pool_)); + factory.Create(std::move(request))); std::shared_ptr store = std::dynamic_pointer_cast(realtime_store); ASSERT_NE(nullptr, store); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 000000000..421eb0c8d --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,309 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/primary_key_realtime_store.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon { + +namespace { + +struct StoredBatch { + std::shared_ptr data; + OffsetRange offset_range; + uint64_t memory_usage; +}; + +class Segment final : public RealtimeSegmentHandle { + public: + Segment(const OffsetRange& range, std::vector&& batches) + : range_(range), batches_(std::move(batches)) {} + + OffsetRange GetOffsetRange() const override { + return range_; + } + const std::vector& Batches() const { + return batches_; + } + + private: + OffsetRange range_; + std::vector batches_; +}; + +class ReadView final : public RealtimeReadView { + public: + explicit ReadView(std::vector>&& segments) + : segments_(std::move(segments)) { + if (!segments_.empty()) { + range_ = OffsetRange(segments_.front()->GetOffsetRange().begin, + segments_.back()->GetOffsetRange().end); + } + } + + std::optional GetOffsetRange() const override { + return range_; + } + const std::vector>& Segments() const { + return segments_; + } + + private: + std::vector> segments_; + std::optional range_; +}; + +class StoredBatchReader final : public BatchReader { + public: + explicit StoredBatchReader(const StoredBatch& batch, + std::shared_ptr arrow_pool) + : arrow_pool_(std::move(arrow_pool)), + data_(batch.data), + metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!data_) { + return MakeEofBatch(); + } + auto array = std::make_unique(); + auto schema = std::make_unique(); + ScopeGuard export_guard([array_ptr = array.get(), schema_ptr = schema.get()]() { + ArrowArrayRelease(array_ptr); + ArrowSchemaRelease(schema_ptr); + }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr record_batch, + arrow::RecordBatch::FromStructArray(data_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr normalized_batch, + ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(array.get(), arrow_pool_)); + data_.reset(); + arrow_pool_.reset(); + export_guard.Release(); + return ReadBatch(std::move(array), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + void Close() override { + data_.reset(); + arrow_pool_.reset(); + } + + private: + std::shared_ptr arrow_pool_; + std::shared_ptr data_; + std::shared_ptr metrics_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + Impl(std::shared_ptr transport_schema, + std::shared_ptr arrow_pool) + : transport_schema_(std::move(transport_schema)), arrow_pool_(std::move(arrow_pool)) {} + + Status Write(RealtimeWriteBatch&& write_batch) { + if (!write_batch.batch || !write_batch.batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t row_count = write_batch.batch->GetData()->length; + if (write_batch.offset_range.begin < 0 || write_batch.offset_range.Count() != row_count || + row_count <= 0) { + return Status::Invalid("PK real-time offset range does not match batch row count"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(write_batch.batch->GetData(), + arrow::struct_(transport_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time transport batch is not a StructArray"); + } + std::shared_ptr transport = + checked_pointer_cast(array); + std::lock_guard lock(mutex_); + building_.push_back(StoredBatch{transport, write_batch.offset_range, + ArrowUtils::GetArrayMemoryUsage(transport->data())}); + building_memory_usage_ += building_.back().memory_usage; + return Status::OK(); + } + + Result>> SealForCommit() { + std::lock_guard lock(mutex_); + if (building_.empty()) { + return std::optional>(); + } + OffsetRange range(building_.front().offset_range.begin, building_.back().offset_range.end); + std::shared_ptr segment = std::make_shared(range, std::move(building_)); + sealed_.push_back(segment); + building_.clear(); + building_memory_usage_ = 0; + return std::optional>(std::move(segment)); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& handle) { + std::shared_ptr segment = std::dynamic_pointer_cast(handle); + if (!segment) { + return Status::Invalid("segment was not created by the PK real-time store"); + } + std::vector> readers; + readers.reserve(segment->Batches().size()); + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back(std::make_unique(batch, arrow_pool_)); + } + return readers; + } + + Result> AcquireReadView() { + std::lock_guard lock(mutex_); + std::vector> segments = sealed_; + if (!building_.empty()) { + OffsetRange range(building_.front().offset_range.begin, + building_.back().offset_range.end); + segments.push_back( + std::make_shared(range, std::vector(building_))); + } + return std::make_shared(std::move(segments)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t, + const RealtimeQueryContext& context) { + std::shared_ptr typed = std::dynamic_pointer_cast(view); + if (!typed) { + return Status::Invalid("read view was not created by the PK real-time store"); + } + if (context.read_schema == nullptr || context.read_schema->release == nullptr) { + return Status::Invalid("PK real-time query read schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, + arrow::ImportSchema(context.read_schema)); + std::vector> readers; + for (const std::shared_ptr& segment : typed->Segments()) { + for (const StoredBatch& batch : segment->Batches()) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr projected, + NestedProjectionUtils::AlignArrayToReadType( + batch.data, arrow::struct_(read_schema->fields()), arrow_pool_.get())); + if (!projected || projected->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + "PK memory query projection did not produce a StructArray"); + } + StoredBatch query_batch{checked_pointer_cast(projected), + batch.offset_range, /*memory_usage=*/0}; + readers.push_back(std::make_unique(query_batch, arrow_pool_)); + } + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_end_offset) { + std::lock_guard lock(mutex_); + auto first_retained = std::find_if( + sealed_.begin(), sealed_.end(), [committed_end_offset](const auto& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + sealed_.erase(sealed_.begin(), first_retained); + return Status::OK(); + } + + uint64_t GetMemoryUsage() const { + std::lock_guard lock(mutex_); + uint64_t total = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_) { + for (const StoredBatch& batch : segment->Batches()) { + total += batch.memory_usage; + } + } + return total; + } + + private: + std::shared_ptr transport_schema_; + std::shared_ptr arrow_pool_; + mutable std::mutex mutex_; + std::vector building_; + std::vector> sealed_; + uint64_t building_memory_usage_ = 0; +}; + +PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) + : impl_(std::move(impl)) {} +PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; + +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& memory_pool) { + if (!memory_pool) { + return Status::Invalid("PK real-time store memory pool is null"); + } + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + return std::shared_ptr(new PrimaryKeyRealtimeStore( + std::make_unique(transport_schema, std::move(arrow_pool)))); +} +Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { + return impl_->Write(std::move(batch)); +} +Result>> +PrimaryKeyRealtimeStore::SealForCommit() { + return impl_->SealForCommit(); +} +Result>> PrimaryKeyRealtimeStore::CreateCommitReaders( + const std::shared_ptr& segment) { + return impl_->CreateCommitReaders(segment); +} +Result> PrimaryKeyRealtimeStore::AcquireReadView() { + return impl_->AcquireReadView(); +} +Result>> PrimaryKeyRealtimeStore::CreateQueryReaders( + const std::shared_ptr& view, int64_t offset, + const RealtimeQueryContext& context) { + return impl_->CreateQueryReaders(view, offset, context); +} +Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_end_offset) { + return impl_->AdvanceCommittedOffset(committed_end_offset); +} +uint64_t PrimaryKeyRealtimeStore::GetMemoryUsage() const { + return impl_->GetMemoryUsage(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h new file mode 100644 index 000000000..46c0fe8f7 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include + +#include "paimon/realtime/realtime_store.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; + +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& memory_pool); + + ~PrimaryKeyRealtimeStore() override; + + Status Write(RealtimeWriteBatch&& batch) override; + Result>> SealForCommit() override; + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override; + Status AdvanceCommittedOffset(int64_t committed_end_offset) override; + uint64_t GetMemoryUsage() const override; + + private: + class Impl; + explicit PrimaryKeyRealtimeStore(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp new file mode 100644 index 000000000..ed80db275 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,406 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/primary_key_realtime_store.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr FieldWithId(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))) + ->WithNullable(nullable); +} + +std::shared_ptr TransportSchema() { + return RealtimePrimaryKeyLayout::CreateSchema( + {FieldWithId("id", arrow::int64(), 0), FieldWithId("value", arrow::utf8(), 1)}); +} + +std::shared_ptr NestedTransportSchema() { + return RealtimePrimaryKeyLayout::CreateSchema( + {FieldWithId("id", arrow::int64(), 0), + FieldWithId("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}), + 1)}); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(TransportSchema()->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +std::unique_ptr MakeSlicedBatch(const std::shared_ptr& schema, + const std::string& json, int64_t offset, + int64_t length) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie() + ->Slice(offset, length); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +void AssertOffsetsZero(const ArrowArray* array) { + ASSERT_NE(nullptr, array); + ASSERT_EQ(0, array->offset); + for (int64_t child = 0; child < array->n_children; ++child) { + AssertOffsetsZero(array->children[child]); + } + if (array->dictionary) { + AssertOffsetsZero(array->dictionary); + } +} + +Result ReadJson(const std::vector>& readers) { + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + batches.push_back(std::move(array)); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches)); + return result->ToString(); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_FALSE(segment.has_value()); + ASSERT_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), + "write batch is null"); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 1, "one"]])"), OffsetRange(0, 0)}), + "offset range does not match batch row count"); + + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[0, 1, 0, 1, "one"], [0, 2, 1, 2, "two"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 3, 2, 3, "three"]])"), OffsetRange(2, 3)})); + + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); + ASSERT_GT(store->GetMemoryUsage(), 0); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[1, 6, 1, 1, "before"], [0, 5, 0, 3, "three"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_EQ( + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 0,\n 2\n ]\n-- " + "child 1 type: int64\n [\n 6,\n 5,\n 7\n ]\n-- child 2 type: int64\n [\n " + "1,\n 0,\n 2\n ]\n-- child 3 type: int64\n [\n 1,\n 3,\n 2\n ]\n-- child " + "4 type: string\n [\n \"before\",\n \"three\",\n \"after\"\n ]", + actual); + for (const std::unique_ptr& reader : readers) { + reader->Close(); + } +} + +void AssertSlicedBatch(BatchReader* reader) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + std::shared_ptr values = checked_pointer_cast(array); + ASSERT_EQ(2, checked_pointer_cast(values->field(3))->Value(0)); + ASSERT_EQ(3, checked_pointer_cast(values->field(3))->Value(1)); + std::shared_ptr nested = + checked_pointer_cast(values->field(4)); + ASSERT_EQ("two", checked_pointer_cast(nested->field(0))->GetString(0)); + ASSERT_EQ("three", checked_pointer_cast(nested->field(0))->GetString(1)); + std::shared_ptr items = + checked_pointer_cast(nested->field(1)); + std::shared_ptr first_items = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(3, first_items->Value(0)); + ASSERT_EQ(4, first_items->Value(1)); + std::shared_ptr second_items = + checked_pointer_cast(items->value_slice(1)); + ASSERT_EQ(5, second_items->Value(0)); + ASSERT_EQ(6, second_items->Value(1)); + ASSERT_OK_AND_ASSIGN(batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { + std::shared_ptr schema = NestedTransportSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + schema, + R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]], [0, 3, 2, 3, ["three", [5, 6]]], [0, 4, 3, 4, ["four", [7, 8]]]])", + 1, 2), + OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + AssertSlicedBatch(readers[0].get()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(readers, store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + AssertSlicedBatch(readers[0].get()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 5, 2, "two"]])"), OffsetRange(5, 6)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 2, 6, 3, "three"]])"), OffsetRange(6, 7)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr retained_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), retained_view->GetOffsetRange()); + + const uint64_t initial_memory_usage = store->GetMemoryUsage(); + ASSERT_GT(initial_memory_usage, 0); + ASSERT_OK(store->AdvanceCommittedOffset(4)); + ASSERT_EQ(initial_memory_usage, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_LT(store->GetMemoryUsage(), initial_memory_usage); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(5, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(6)); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(6, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(7)); + ASSERT_EQ(0, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_FALSE(current_view->GetOffsetRange().has_value()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(retained_view, /*offset_begin=*/0, context)); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); + ASSERT_NE(std::string::npos, actual.find("\"three\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 2, 1, 1, "one"]])"), OffsetRange(1, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{/*read_schema=*/c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryPoolOutlivesStoreReaderAndExport) { + const std::shared_ptr stored_schema = TransportSchema(); + std::shared_ptr pool = GetMemoryPool(); + std::weak_ptr pool_lifetime = pool; + auto write_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, write_schema.get()).ok()); + ArrowRealtimeStoreFactory factory; + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + factory.Create(RealtimeStoreCreateRequest{ + std::move(write_schema), + /*options=*/{}, pool, RealtimeStoreMode::PRIMARY_KEY})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 7, "seven"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + view.reset(); + store.reset(); + pool.reset(); + ASSERT_FALSE(pool_lifetime.expired()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr imported = std::move(import_result).ValueOrDie(); + readers.clear(); + ASSERT_FALSE(pool_lifetime.expired()); + imported.reset(); + ASSERT_TRUE(pool_lifetime.expired()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { + const std::shared_ptr stored_profile_a = + FieldWithId("profile_a", arrow::int32(), 30); + const std::shared_ptr stored_a = FieldWithId("a", arrow::int32(), 10); + const std::shared_ptr stored_b = FieldWithId("b", arrow::int32(), 11); + const std::shared_ptr stored_x = FieldWithId("x", arrow::int32(), 20); + const std::shared_ptr stored_y = FieldWithId("y", arrow::int32(), 21); + arrow::FieldVector stored_value_fields = { + FieldWithId("id", arrow::int64(), 0), + FieldWithId("profile", arrow::struct_({stored_profile_a}), 1), + FieldWithId("items", arrow::list(arrow::struct_({stored_a, stored_b})), 2), + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_x, stored_y})), 3)}; + std::shared_ptr stored_schema = + RealtimePrimaryKeyLayout::CreateSchema(stored_value_fields); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(stored_schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + stored_schema, + R"([[0, 1, 0, 6, [5], [[1, 2]], [["before", [3, 4]]]], [0, 2, 1, 7, [50], [[100, 200], null], [["k1", [7, 8]], ["k2", null]]], [0, 3, 2, 8, [500], [[9, 10]], [["after", [11, 12]]]]])", + 1, 1), + OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + arrow::FieldVector requested_value_fields; + requested_value_fields.push_back(FieldWithId("profile", arrow::struct_({stored_profile_a}), 1)); + requested_value_fields.push_back( + FieldWithId("items", arrow::list(arrow::struct_({stored_b, stored_a})), 2)); + requested_value_fields.push_back( + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_y, stored_x})), 3)); + std::shared_ptr requested_schema = + RealtimePrimaryKeyLayout::CreateSchema(requested_value_fields); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*requested_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(requested_schema->fields()))); + std::shared_ptr projected = checked_pointer_cast(array); + const std::shared_ptr profile = + checked_pointer_cast(projected->field(3)); + ASSERT_EQ(50, checked_pointer_cast(profile->field(0))->Value(0)); + const std::shared_ptr items = + checked_pointer_cast(projected->field(4)); + const std::shared_ptr item_values = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(200, checked_pointer_cast(item_values->field(0))->Value(0)); + ASSERT_EQ(100, checked_pointer_cast(item_values->field(1))->Value(0)); + ASSERT_TRUE(item_values->IsNull(1)); + + const std::shared_ptr attrs = + checked_pointer_cast(projected->field(5)); + const int64_t attr_offset = attrs->value_offset(0); + const int64_t attr_length = attrs->value_length(0); + const std::shared_ptr attr_keys = + checked_pointer_cast(attrs->keys()->Slice(attr_offset, attr_length)); + ASSERT_EQ("k1", attr_keys->GetString(0)); + const std::shared_ptr attr_values = + checked_pointer_cast(attrs->items()->Slice(attr_offset, attr_length)); + ASSERT_EQ(8, checked_pointer_cast(attr_values->field(0))->Value(0)); + ASSERT_EQ(7, checked_pointer_cast(attr_values->field(1))->Value(0)); + ASSERT_TRUE(attr_values->IsNull(1)); +} + +} // namespace +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 9d519d791..632d64e16 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -55,10 +55,11 @@ Result> RealtimeAppendOnlyWriter::Crea } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, RealtimeContextImpl::Cast(realtime_context)); - PAIMON_ASSIGN_OR_RAISE( - RealtimeStoreState store_state, - realtime_context_impl->GetOrCreateRealtimeStore(partition, bucket, std::move(write_schema), - statistics_mode, options, memory_pool)); + RealtimeStoreCreateRequest request{std::move(write_schema), options, memory_pool, + RealtimeStoreMode::APPEND_ONLY, statistics_mode}; + PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore( + std::move(request), RealtimePartitionBucket(partition, bucket))); return std::shared_ptr(new RealtimeAppendOnlyWriter( store_state.store, file_writer, input_schema, store_state.initial_offset, memory_pool)); } diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index f6bad5cf1..2b9d7d07c 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -34,14 +34,33 @@ #include #include +#include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "fmt/format.h" #include "paimon/arrow/abi.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/uuid.h" #include "paimon/macros.h" #include "paimon/realtime/realtime_store.h" #include "paimon/status.h" namespace paimon { +namespace { + +std::string PartitionToString(const std::map& partition) { + std::string result = "{"; + for (auto iter = partition.begin(); iter != partition.end(); ++iter) { + if (iter != partition.begin()) { + result += ", "; + } + result += iter->first + "=" + iter->second; + } + return result + "}"; +} + +} // namespace RealtimeContextImpl::RealtimeContextImpl(const std::shared_ptr& factory) : factory_(factory) {} @@ -78,31 +97,36 @@ Status RealtimeContextImpl::Start() { } Result RealtimeContextImpl::GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) { + RealtimeStoreCreateRequest&& request, const RealtimePartitionBucket& partition_bucket) { + if (!request.write_schema || !request.write_schema->release) { + return Status::Invalid("real-time store write schema is null"); + } + ScopeGuard schema_guard( + [schema = request.write_schema.get()]() { ArrowSchemaRelease(schema); }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr requested_schema, + arrow::ImportSchema(request.write_schema.get())); + schema_guard.Release(); std::lock_guard progress_lock(progress_mutex_); std::lock_guard registry_lock(mutex_); - const RealtimePartitionBucket key(partition, bucket); + auto iter = stores_.find(partition_bucket); int64_t initial_offset = 0; - auto offset_iter = committed_offsets_.find(key); + auto offset_iter = committed_offsets_.find(partition_bucket); if (offset_iter != committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); - } return Status::Invalid("real-time offset has reached INT64_MAX"); } initial_offset = offset_iter->second; } - auto iter = stores_.find(key); if (iter != stores_.end()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); + if (iter->second.mode != request.mode || + !iter->second.write_schema->Equals(*requested_schema, /*check_metadata=*/true)) { + return Status::Invalid(fmt::format( + "real-time store schema or mode mismatch for partition {}, bucket {}; recreate " + "the RealtimeContext", + PartitionToString(partition_bucket.partition), partition_bucket.bucket)); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - iter->second->AcquireReadView()); + iter->second.store->AcquireReadView()); if (!read_view) { return Status::Invalid("real-time store returned a null read view"); } @@ -117,27 +141,54 @@ Result RealtimeContextImpl::GetOrCreateRealtimeStore( initial_offset = memory_range->end; } } - return RealtimeStoreState{iter->second, initial_offset}; + return RealtimeStoreState{iter->second.store, initial_offset}; + } + if (!request.memory_pool) { + return Status::Invalid("real-time store memory pool is null"); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*requested_schema, request.write_schema.get())); + RealtimeStoreMode mode = request.mode; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, + factory_->Create(std::move(request))); + if (!store) { + return Status::Invalid("real-time store factory returned a null store"); } - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr store, - factory_->Create(std::move(write_schema), statistics_mode, options, memory_pool)); - stores_.emplace(key, store); + stores_.emplace(partition_bucket, StoreEntry{store, requested_schema, mode}); if (offset_iter != committed_offsets_.end()) { - reclaimed_offsets_.emplace(key, offset_iter->second); + reclaimed_offsets_.emplace(partition_bucket, offset_iter->second); } return RealtimeStoreState{std::move(store), initial_offset}; } +Result RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + std::lock_guard lock(mutex_); + auto iter = stores_.find(partition_bucket); + if (iter == stores_.end()) { + return Status::KeyError(fmt::format("real-time store not found for partition {}, bucket {}", + PartitionToString(partition_bucket.partition), + partition_bucket.bucket)); + } + StoreEntry& entry = iter->second; + if (max_sequence_number > entry.materialized_max_sequence_number) { + entry.materialized_max_sequence_number = max_sequence_number; + } + return entry.materialized_max_sequence_number; +} + Result> RealtimeContextImpl::AcquireReadViews() { std::lock_guard lock(mutex_); std::vector result; result.reserve(stores_.size()); for (const auto& [partition_bucket, store] : stores_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - store->AcquireReadView()); + store.store->AcquireReadView()); + if (!read_view) { + return Status::Invalid("real-time store returned a null read view"); + } result.push_back( - RealtimePartitionBucketView{partition_bucket, store, std::move(read_view)}); + RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; } @@ -270,7 +321,7 @@ Status RealtimeContextImpl::AdvanceCommittedProgress(int64_t snapshot_id, } auto store_iter = stores_.find(partition_bucket); if (store_iter != stores_.end()) { - notifications.emplace_back(partition_bucket, store_iter->second, + notifications.emplace_back(partition_bucket, store_iter->second.store, committed_end_offset); } } diff --git a/src/paimon/core/realtime/realtime_context_impl.h b/src/paimon/core/realtime/realtime_context_impl.h index 66c324cab..ea069a5cd 100644 --- a/src/paimon/core/realtime/realtime_context_impl.h +++ b/src/paimon/core/realtime/realtime_context_impl.h @@ -32,12 +32,16 @@ #include #include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" #include "paimon/result.h" -#include "paimon/statistics_mode.h" #include "paimon/visibility.h" struct ArrowSchema; +namespace arrow { +class Schema; +} // namespace arrow + namespace paimon { class RealtimeStore; @@ -66,10 +70,10 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { const std::shared_ptr& context); Result GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool); + RealtimeStoreCreateRequest&& request, const RealtimePartitionBucket& partition_bucket); + + Result AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number); Result> AcquireReadViews(); @@ -93,6 +97,13 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::chrono::steady_clock::time_point expire_at; }; + struct StoreEntry { + std::shared_ptr store; + std::shared_ptr write_schema; + RealtimeStoreMode mode; + int64_t materialized_max_sequence_number = -1; + }; + explicit RealtimeContextImpl(const std::shared_ptr& factory); Status Start(); @@ -102,7 +113,7 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::shared_ptr factory_; std::mutex mutex_; std::mutex progress_mutex_; - std::map> stores_; + std::map stores_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; // Progress already reflected in stores owned by this context. diff --git a/src/paimon/core/realtime/realtime_context_test.cpp b/src/paimon/core/realtime/realtime_context_test.cpp index 017820fd4..c9531eda5 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -9,12 +9,11 @@ * * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ #include @@ -31,7 +30,6 @@ #include "arrow/c/helpers.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/memory/memory_pool.h" -#include "paimon/realtime/realtime_store.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -49,26 +47,24 @@ class TestingRealtimeStore : public RealtimeStore { Status Write(RealtimeWriteBatch&&) override { return Status::OK(); } - Result>> SealForCommit() override { return std::optional>(); } - Result>> CreateCommitReaders( const std::shared_ptr&) override { return std::vector>(); } - Result> AcquireReadView() override { ++acquire_count; + if (return_null_read_view) { + return std::shared_ptr(); + } return std::make_shared(); } - Result>> CreateQueryReaders( const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { return std::vector>(); } - Status AdvanceCommittedOffset(int64_t committed_offset) override { ++advance_count; if (fail_next_advance) { @@ -78,7 +74,6 @@ class TestingRealtimeStore : public RealtimeStore { committed_offsets.push_back(committed_offset); return Status::OK(); } - uint64_t GetMemoryUsage() const override { return 0; } @@ -86,33 +81,37 @@ class TestingRealtimeStore : public RealtimeStore { int32_t acquire_count = 0; int32_t advance_count = 0; bool fail_next_advance = false; + bool return_null_read_view = false; std::vector committed_offsets; }; class TestingRealtimeStoreFactory : public RealtimeStoreFactory { public: - Result> Create(std::unique_ptr write_schema, - StatisticsMode, - const std::map&, - const std::shared_ptr&) override { - if (!write_schema || !write_schema->release) { + Result> Create(RealtimeStoreCreateRequest&& request) override { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("testing write schema is null"); } - ArrowSchemaRelease(write_schema.get()); + ArrowSchemaRelease(request.write_schema.get()); + if (return_null_store) { + return std::shared_ptr(); + } auto store = std::make_shared(); stores.push_back(store); return store; } + bool return_null_store = false; std::vector> stores; }; -std::unique_ptr MakeWriteSchema() { - auto c_schema = std::make_unique(); +std::unique_ptr MakeWriteSchema( + const std::shared_ptr& id_type = arrow::int64(), + const std::shared_ptr& metadata = nullptr) { + auto schema = std::make_unique(); EXPECT_TRUE( - arrow::ExportSchema(*arrow::schema({arrow::field("id", arrow::int64())}), c_schema.get()) + arrow::ExportSchema(*arrow::schema({arrow::field("id", id_type)}, metadata), schema.get()) .ok()); - return c_schema; + return schema; } Result> CreateContext( @@ -122,35 +121,42 @@ Result> CreateContext( return RealtimeContextImpl::Cast(context); } +Result GetOrCreateAppendStore( + const std::shared_ptr& context, + const std::map& partition, int32_t bucket, + std::unique_ptr write_schema, const std::map& options, + const std::shared_ptr& memory_pool, + StatisticsMode statistics_mode = StatisticsMode::NONE) { + return context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(write_schema), options, memory_pool, + RealtimeStoreMode::APPEND_ONLY, statistics_mode}, + RealtimePartitionBucket(partition, bucket)); +} + TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - + ASSERT_OK_AND_ASSIGN(RealtimeStoreState first, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, + MakeWriteSchema(), {{"k", "v"}}, GetDefaultPool())); + ASSERT_EQ(0, first.initial_offset); ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {{"k", "v"}}, pool)); - ASSERT_EQ(0, first_state.initial_offset); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_again_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_EQ(first_state.store, first_again_state.store); - ASSERT_EQ(0, first_again_state.initial_offset); + RealtimeStoreState second, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, MakeWriteSchema(), {}, + GetDefaultPool(), StatisticsMode::FULL)); + ASSERT_EQ(first.store, second.store); + ASSERT_EQ(0, second.initial_offset); ASSERT_EQ(1, factory->stores.size()); ASSERT_EQ(1, factory->stores[0]->acquire_count); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState second_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-02"}}, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState third_state, - context->GetOrCreateRealtimeStore({{"dt", "2026-08-03"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_NE(first_state.store, second_state.store); - ASSERT_NE(first_state.store, third_state.store); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState third, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 1, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState fourth, + GetOrCreateAppendStore(context, {{"dt", "2026-08-03"}}, 0, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_NE(first.store, third.store); + ASSERT_NE(first.store, fourth.store); ASSERT_EQ(3, factory->stores.size()); ASSERT_OK_AND_ASSIGN(std::vector views, @@ -158,23 +164,87 @@ TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { ASSERT_EQ(3, views.size()); const RealtimePartitionBucket expected_partition_bucket({{"dt", "2026-08-02"}}, 0); ASSERT_EQ(expected_partition_bucket, views[0].partition_bucket); - ASSERT_EQ(first_state.store, views[0].store); + ASSERT_EQ(first.store, views[0].store); ASSERT_TRUE(views[0].read_view); ASSERT_EQ(2, factory->stores[0]->acquire_count); ASSERT_EQ(1, factory->stores[1]->acquire_count); ASSERT_EQ(1, factory->stores[2]->acquire_count); } +TEST(RealtimeContextTest, TestRejectsMismatchedModeOnStoreReuse) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + + ASSERT_NOK_WITH_MSG( + context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{ + MakeWriteSchema(), {}, GetDefaultPool(), RealtimeStoreMode::PRIMARY_KEY}, + RealtimePartitionBucket(partition, 0)), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_EQ(1, factory->stores.size()); +} + +TEST(RealtimeContextTest, TestRejectsMismatchedSchemaOnStoreReuse) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + std::shared_ptr metadata = + arrow::key_value_metadata({"identity"}, {"v1"}); + ASSERT_OK(GetOrCreateAppendStore( + context, partition, 0, MakeWriteSchema(arrow::int64(), metadata), {}, GetDefaultPool())); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(arrow::int32(), metadata), {}, + GetDefaultPool()), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_NOK_WITH_MSG( + GetOrCreateAppendStore( + context, partition, 0, + MakeWriteSchema(arrow::int64(), arrow::key_value_metadata({"identity"}, {"v2"})), {}, + GetDefaultPool()), + "schema or mode mismatch for partition {dt=2026-08-02}, bucket 0; recreate the " + "RealtimeContext"); + ASSERT_EQ(1, factory->stores.size()); +} + +TEST(RealtimeContextTest, TestReconcilesPrimaryKeyInitialSequence) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + const std::map partition = {{"dt", "2026-08-02"}}; + const RealtimePartitionBucket partition_bucket(partition, /*bucket=*/0); + + ASSERT_OK(context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{MakeWriteSchema(), /*options=*/{}, GetDefaultPool(), + RealtimeStoreMode::PRIMARY_KEY}, + partition_bucket)); + + ASSERT_OK_AND_ASSIGN(int64_t first, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/4)); + ASSERT_EQ(4, first); + ASSERT_OK_AND_ASSIGN(int64_t second, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/8)); + ASSERT_EQ(8, second); + ASSERT_OK_AND_ASSIGN(int64_t third, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/6)); + ASSERT_EQ(8, third); + ASSERT_OK_AND_ASSIGN(int64_t fourth, context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, /*max_sequence_number=*/10)); + ASSERT_EQ(10, fourth); +} + TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(2, factory->stores.size()); ASSERT_NOK_WITH_MSG(context->AdvanceCommittedProgress(-1, {}), @@ -191,10 +261,9 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); ASSERT_TRUE(factory->stores[1]->committed_offsets.empty()); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState restored_state, - context->GetOrCreateRealtimeStore({{"dt", "unknown"}}, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK_AND_ASSIGN(RealtimeStoreState restored_state, + GetOrCreateAppendStore(context, {{"dt", "unknown"}}, 0, MakeWriteSchema(), + {}, GetDefaultPool())); ASSERT_EQ(9, restored_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress( @@ -214,7 +283,6 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { TEST(RealtimeContextTest, TestRemovedInactivePartitionDoesNotRequireReopen) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map active_partition = {{"dt", "2026-08-02"}}; const std::map inactive_partition = {{"dt", "2026-08-03"}}; const RealtimePartitionBucket active_partition_bucket(active_partition, /*bucket=*/0); @@ -223,29 +291,28 @@ TEST(RealtimeContextTest, TestRemovedInactivePartitionDoesNotRequireReopen) { ASSERT_OK(context->AdvanceCommittedProgress( 5, {{active_partition_bucket, /*offset=*/7}, {inactive_partition_bucket, /*offset=*/9}})); ASSERT_OK_AND_ASSIGN(RealtimeStoreState active_state, - context->GetOrCreateRealtimeStore(active_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, active_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_EQ(7, active_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress(6, {{active_partition_bucket, /*offset=*/7}})); ASSERT_OK_AND_ASSIGN(RealtimeStoreState inactive_state, - context->GetOrCreateRealtimeStore(inactive_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + GetOrCreateAppendStore(context, inactive_partition, 0, MakeWriteSchema(), + {}, GetDefaultPool())); ASSERT_EQ(0, inactive_state.initial_offset); } TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 2, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 2, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(3, factory->stores.size()); factory->stores[1]->fail_next_advance = true; @@ -259,9 +326,9 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_TRUE(factory->stores[1]->committed_offsets.empty()); ASSERT_EQ(std::vector({9}), factory->stores[2]->committed_offsets); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState failed_store_state, - context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK_AND_ASSIGN( + RealtimeStoreState failed_store_state, + GetOrCreateAppendStore(context, partition, 1, MakeWriteSchema(), {}, GetDefaultPool())); ASSERT_EQ(8, failed_store_state.initial_offset); ASSERT_OK(context->AdvanceCommittedProgress(5, committed_offsets)); @@ -274,16 +341,15 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { TEST(RealtimeContextTest, TestRequiresReopenWhenCommittedProgressMovesBackwards) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map first_partition = {{"dt", "2026-08-02"}}; const std::map second_partition = {{"dt", "2026-08-03"}}; const RealtimePartitionBucket first_partition_bucket(first_partition, /*bucket=*/0); const RealtimePartitionBucket second_partition_bucket(second_partition, /*bucket=*/0); - ASSERT_OK(context->GetOrCreateRealtimeStore(first_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(second_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK(GetOrCreateAppendStore(context, first_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, second_partition, 0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK(context->AdvanceCommittedProgress( 5, {{first_partition_bucket, /*offset=*/7}, {second_partition_bucket, /*offset=*/9}})); ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); @@ -308,8 +374,8 @@ TEST(RealtimeContextTest, TestRequiresReopenWhenCommittedProgressMovesBackwards) TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); @@ -332,8 +398,8 @@ TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { TEST(RealtimeContextTest, TestExpiresAbandonedReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); @@ -354,5 +420,20 @@ TEST(RealtimeContextTest, TestRejectsNullFactory) { "real-time store factory is null"); } +TEST(RealtimeContextTest, TestRejectsNullPluginResults) { + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); + factory->return_null_store = true; + ASSERT_NOK_WITH_MSG(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, + MakeWriteSchema(), {}, GetDefaultPool()), + "real-time store factory returned a null store"); + + factory->return_null_store = false; + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); + factory->stores[0]->return_null_read_view = true; + ASSERT_NOK_WITH_MSG(context->AcquireReadViews(), "real-time store returned a null read view"); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.cpp b/src/paimon/core/realtime/realtime_primary_key_reader.cpp new file mode 100644 index 000000000..19ba0a985 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader.cpp @@ -0,0 +1,510 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/realtime_primary_key_reader.h" + +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_primitive.h" +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "fmt/format.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/status.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { + +namespace { + +template +void CloseReaders(const std::vector>& readers) { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + +class RealtimeOffsetCoverage { + public: + static Result> Create(const OffsetRange& offsets, + size_t reader_count, + bool allow_committed_prefix) { + if (offsets.begin < 0 || offsets.end < offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid offset range"); + } + return std::shared_ptr( + new RealtimeOffsetCoverage(offsets, reader_count, allow_committed_prefix)); + } + + Status Add(const arrow::Int64Array& offsets) { + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (allow_committed_prefix_ && offset < 0) { + return Status::Invalid("PK real-time store reader offset must be non-negative"); + } + if (allow_committed_prefix_ && offset < offsets_.begin) { + continue; + } + if (offset < offsets_.begin || offset >= offsets_.end) { + return Status::Invalid( + allow_committed_prefix_ + ? "PK real-time store query reader offset is outside the visible range" + : "PK real-time store commit reader offset is outside the sealed range"); + } + if (!seen_offsets_.CheckedAdd(offset)) { + return CoverageError(); + } + } + return Status::OK(); + } + + Status FinishReader() { + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && + seen_offsets_.Cardinality() != offsets_.Count()) { + return CoverageError(); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& offsets, size_t reader_count, + bool allow_committed_prefix) + : offsets_(offsets), + reader_count_(reader_count), + allow_committed_prefix_(allow_committed_prefix) {} + + Status CoverageError() const { + return Status::Invalid( + allow_committed_prefix_ + ? "PK real-time store query readers did not cover the visible range" + : "PK real-time store commit readers did not cover the sealed range"); + } + + OffsetRange offsets_; + size_t reader_count_; + bool allow_committed_prefix_; + RoaringBitmap64 seen_offsets_; + size_t finished_reader_count_ = 0; +}; + +Status CheckTransportField(const std::shared_ptr& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid( + fmt::format("realtime primary-key transport schema is missing field {} at index {}", + expected_field.Name(), field_idx)); + } + const std::shared_ptr& field = schema->field(field_idx); + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field->name() != expected_field.Name() || !field->type()->Equals(*expected_field.Type()) || + field->nullable() || field_id != expected_field.Id()) { + return Status::Invalid(fmt::format( + "realtime primary-key transport schema field {} must be non-null {}:{} with field id " + "{}, got {}:{} nullable={} field id {}", + field_idx, expected_field.Name(), expected_field.Type()->ToString(), + expected_field.Id(), field->name(), field->type()->ToString(), field->nullable(), + field_id)); + } + return Status::OK(); +} + +Result> ResolveFieldIndexes( + const std::shared_ptr& transport_schema, + const std::unordered_map& field_indexes, + const std::shared_ptr& row_schema) { + std::vector result; + result.reserve(row_schema->num_fields()); + for (const std::shared_ptr& row_field : row_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(row_field)); + auto field_index = field_indexes.find(field_id); + if (field_index == field_indexes.end()) { + return Status::Invalid(fmt::format( + "cannot find field id {} in realtime primary-key transport schema", field_id)); + } + const std::shared_ptr& transport_field = + transport_schema->field(field_index->second); + if (!transport_field->type()->Equals(row_field->type())) { + return Status::Invalid(fmt::format( + "realtime primary-key transport field id {} type {} does not match row type {}", + field_id, transport_field->type()->ToString(), row_field->type()->ToString())); + } + result.push_back(field_index->second); + } + return result; +} + +class RealtimePrimaryKeyReaderPlan { + public: + static Result> Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema) { + std::unordered_map field_indexes; + field_indexes.reserve(transport_schema->num_fields() - + RealtimePrimaryKeyLayout::kValueStartIndex); + for (int32_t i = RealtimePrimaryKeyLayout::kValueStartIndex; + i < transport_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId( + transport_schema->field(i))); + if (!field_indexes.emplace(field_id, i).second) { + return Status::Invalid(fmt::format( + "duplicate field id {} in realtime primary-key transport schema", field_id)); + } + } + PAIMON_ASSIGN_OR_RAISE(std::vector key_field_indexes, + ResolveFieldIndexes(transport_schema, field_indexes, key_schema)); + PAIMON_ASSIGN_OR_RAISE(std::vector value_field_indexes, + ResolveFieldIndexes(transport_schema, field_indexes, value_schema)); + return std::shared_ptr(new RealtimePrimaryKeyReaderPlan( + transport_schema, std::move(key_field_indexes), std::move(value_field_indexes))); + } + + const std::shared_ptr& TransportSchema() const { + return transport_schema_; + } + + const std::vector& KeyFieldIndexes() const { + return key_field_indexes_; + } + + const std::vector& ValueFieldIndexes() const { + return value_field_indexes_; + } + + private: + RealtimePrimaryKeyReaderPlan(const std::shared_ptr& schema, + std::vector&& key_indexes, + std::vector&& value_indexes) + : transport_schema_(schema), + key_field_indexes_(std::move(key_indexes)), + value_field_indexes_(std::move(value_indexes)) {} + + const std::shared_ptr transport_schema_; + const std::vector key_field_indexes_; + const std::vector value_field_indexes_; +}; + +class RealtimePrimaryKeyReader final : public KeyValueRecordReader { + public: + RealtimePrimaryKeyReader(std::unique_ptr&& reader, + const std::shared_ptr& plan, + const std::optional& visible_offsets, + const std::shared_ptr& pool, + const std::shared_ptr& offset_coverage) + : reader_(std::move(reader)), + plan_(plan), + visible_offsets_(visible_offsets), + pool_(pool), + offset_coverage_(offset_coverage) {} + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(RealtimePrimaryKeyReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->RowCount(); + } + + Result Next() override { + if (cursor_ >= reader_->RowCount()) { + return Status::Invalid("No more realtime primary-key values in current iterator"); + } + const int64_t row = reader_->RowAt(cursor_); + std::shared_ptr key = + std::make_shared(reader_->key_ctx_, row); + auto value = std::make_unique(reader_->value_ctx_, row); + PAIMON_ASSIGN_OR_RAISE(const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kind_array_->Value(row))); + int64_t sequence_number = reader_->sequence_number_array_->Value(row); + ++cursor_; + return KeyValue(row_kind, sequence_number, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + RealtimePrimaryKeyReader* reader_; + int64_t cursor_ = 0; + }; + + Result> NextBatch() override { + return NextBatchImpl(); + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + ResetBatchState(); + reader_->Close(); + } + + private: + Result> NextBatchImpl() { + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } + return std::unique_ptr(); + } + auto& [batch, selection] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + "cannot cast realtime primary-key transport batch to StructArray"); + } + std::shared_ptr data_batch = + checked_pointer_cast(arrow_array); + PAIMON_RETURN_NOT_OK(ValidateTransportBatch(data_batch)); + + std::shared_ptr> offset_array = + checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex)); + if (offset_coverage_) { + PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array)); + } + + row_kind_array_ = checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kValueKindIndex)); + sequence_number_array_ = checked_pointer_cast>( + data_batch->field(RealtimePrimaryKeyLayout::kSequenceNumberIndex)); + arrow::ArrayVector key_fields; + key_fields.reserve(plan_->KeyFieldIndexes().size()); + for (int32_t index : plan_->KeyFieldIndexes()) { + key_fields.push_back(data_batch->field(index)); + } + arrow::ArrayVector value_fields; + value_fields.reserve(plan_->ValueFieldIndexes().size()); + for (int32_t index : plan_->ValueFieldIndexes()) { + value_fields.push_back(data_batch->field(index)); + } + key_ctx_ = std::make_shared(key_fields, pool_); + value_ctx_ = std::make_shared(value_fields, pool_); + PAIMON_ASSIGN_OR_RAISE(bool has_selected_rows, + SelectRows(*offset_array, std::move(selection))); + if (!has_selected_rows) { + continue; + } + ArrowUtils::TraverseArray(data_batch); + return std::make_unique(this); + } + } + + Status ValidateTransportBatch(const std::shared_ptr& data_batch) const { + if (data_batch->num_fields() != plan_->TransportSchema()->num_fields()) { + return Status::Invalid(fmt::format( + "realtime primary-key transport batch field count {} does not match schema field " + "count {}", + data_batch->num_fields(), plan_->TransportSchema()->num_fields())); + } + const arrow::FieldVector& batch_fields = data_batch->type()->fields(); + for (int32_t i = 0; i < data_batch->num_fields(); ++i) { + if (!batch_fields[i]->Equals(plan_->TransportSchema()->field(i), true)) { + return Status::Invalid(fmt::format( + "realtime primary-key transport batch field {} does not match declared schema", + i)); + } + } + if (data_batch->field(RealtimePrimaryKeyLayout::kValueKindIndex)->null_count() != 0 || + data_batch->field(RealtimePrimaryKeyLayout::kSequenceNumberIndex)->null_count() != 0 || + data_batch->field(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex)->null_count() != 0) { + return Status::Invalid("realtime primary-key transport columns must not contain nulls"); + } + return Status::OK(); + } + + Result SelectRows(const arrow::Int64Array& offsets, RoaringBitmap32&& selection) { + for (auto iter = selection.Begin(); iter != selection.End(); ++iter) { + const uint32_t row = *iter; + if (static_cast(row) >= offsets.length()) { + return Status::Invalid( + fmt::format("selected row id {} is out of bounds for realtime primary-key " + "transport batch length {}", + row, offsets.length())); + } + } + if (selection.Cardinality() != offsets.length()) { + return Status::Invalid( + "PK real-time store reader bitmap must cover every raw " + "transport row"); + } + selected_rows_.reserve(offsets.length()); + for (int64_t row = 0; row < offsets.length(); ++row) { + if (!visible_offsets_.has_value() || (offsets.Value(row) >= visible_offsets_->begin && + offsets.Value(row) < visible_offsets_->end)) { + selected_rows_.push_back(row); + } + } + return !selected_rows_.empty(); + } + + int64_t RowCount() const { + return static_cast(selected_rows_.size()); + } + + int64_t RowAt(int64_t ordinal) const { + return selected_rows_[ordinal]; + } + + void ResetBatchState() { + key_ctx_.reset(); + value_ctx_.reset(); + row_kind_array_.reset(); + sequence_number_array_.reset(); + selected_rows_.clear(); + } + + private: + std::unique_ptr reader_; + std::shared_ptr plan_; + std::optional visible_offsets_; + std::shared_ptr pool_; + std::shared_ptr offset_coverage_; + bool offset_coverage_finished_ = false; + std::shared_ptr key_ctx_; + std::shared_ptr value_ctx_; + std::shared_ptr> row_kind_array_; + std::shared_ptr> sequence_number_array_; + std::vector selected_rows_; +}; + +} // namespace + +std::shared_ptr RealtimePrimaryKeyLayout::CreateSchema( + const std::vector>& value_fields) { + arrow::FieldVector fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())}; + fields.insert(fields.end(), value_fields.begin(), value_fields.end()); + return arrow::schema(std::move(fields)); +} + +Status RealtimePrimaryKeyLayout::ValidateSchema( + const std::shared_ptr& transport_schema) { + if (!transport_schema || transport_schema->num_fields() < kValueStartIndex) { + return Status::Invalid( + "realtime primary-key transport schema must contain transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckTransportField(transport_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK(CheckTransportField(transport_schema, kSequenceNumberIndex, + SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK(CheckTransportField(transport_schema, kRealtimeOffsetIndex, + SpecialFields::RealtimeOffset())); + return Status::OK(); +} + +Result>> +RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard remaining_raw_readers_guard([&readers]() { CloseReaders(readers); }); + if (readers.empty() && visible_offsets.begin < visible_offsets.end) { + return Status::Invalid( + "PK real-time store returned no query readers for a non-empty visible range"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null query reader"); + } + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr plan, + RealtimePrimaryKeyReaderPlan::Create(transport_schema, key_schema, value_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(visible_offsets, readers.size(), + /*allow_committed_prefix=*/true)); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + adapted_readers.push_back(std::make_unique( + std::move(reader), plan, visible_offsets, memory_pool, offset_coverage)); + } + remaining_raw_readers_guard.Release(); + return adapted_readers; +} + +Result>> +RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard remaining_raw_readers_guard([&readers]() { CloseReaders(readers); }); + if (readers.empty()) { + return Status::Invalid( + "PK real-time store returned no commit readers for a sealed segment"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr plan, + RealtimePrimaryKeyReaderPlan::Create(transport_schema, key_schema, value_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(sealed_offsets, readers.size(), + /*allow_committed_prefix=*/false)); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + adapted_readers.push_back(std::make_unique( + std::move(reader), plan, std::nullopt, memory_pool, offset_coverage)); + } + remaining_raw_readers_guard.Release(); + return adapted_readers; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_reader.h b/src/paimon/core/realtime/realtime_primary_key_reader.h new file mode 100644 index 000000000..d175c3b63 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader.h @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "arrow/type_fwd.h" +#include "paimon/core/io/key_value_record_reader.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/result.h" + +namespace paimon { +class BatchReader; +class MemoryPool; + +/// Defines the Arrow field layout for PK realtime transport batches. +class RealtimePrimaryKeyLayout { + public: + RealtimePrimaryKeyLayout() = delete; + ~RealtimePrimaryKeyLayout() = delete; + + static constexpr int32_t kValueKindIndex = 0; + static constexpr int32_t kSequenceNumberIndex = 1; + static constexpr int32_t kRealtimeOffsetIndex = 2; + static constexpr int32_t kValueStartIndex = 3; + + static std::shared_ptr CreateSchema( + const std::vector>& value_fields); + + static Status ValidateSchema(const std::shared_ptr& transport_schema); +}; + +class RealtimePrimaryKeyReaderFactory { + public: + RealtimePrimaryKeyReaderFactory() = delete; + ~RealtimePrimaryKeyReaderFactory() = delete; + + static Result>> CreateForQuery( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + + static Result>> CreateForCommit( + std::vector>&& readers, + const std::shared_ptr& transport_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp new file mode 100644 index 000000000..89a39cd63 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_reader_test.cpp @@ -0,0 +1,688 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/realtime_primary_key_reader.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/array_nested.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/utils/key_value_checker.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +namespace { + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))); +} + +std::shared_ptr MakeTransportSchema(const arrow::FieldVector& value_fields) { + return RealtimePrimaryKeyLayout::CreateSchema(value_fields); +} + +Result> CreateRealtimePrimaryKeyQueryReaderForTest( + std::unique_ptr&& reader, const std::shared_ptr& transport_schema, + const OffsetRange& visible_offsets, const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> readers; + readers.push_back(std::move(reader)); + PAIMON_ASSIGN_OR_RAISE(std::vector> adapted_readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(readers), transport_schema, visible_offsets, key_schema, + value_schema, memory_pool)); + return std::move(adapted_readers[0]); +} + +Result> CreateRealtimePrimaryKeyCommitReaderForTest( + std::unique_ptr&& reader, const std::shared_ptr& transport_schema, + const OffsetRange& sealed_offsets, const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + std::vector> readers; + readers.push_back(std::move(reader)); + PAIMON_ASSIGN_OR_RAISE(std::vector> adapted_readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(readers), transport_schema, sealed_offsets, key_schema, + value_schema, memory_pool)); + return std::move(adapted_readers[0]); +} + +class TrackingBatchReader : public BatchReader { + public: + TrackingBatchReader(std::unique_ptr&& delegate, int32_t* close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + ++(*close_count_); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + int32_t* close_count_; +}; + +class MalformedBitmapBatchReader : public BatchReader { + public: + MalformedBitmapBatchReader(std::unique_ptr&& delegate, int32_t row_id) + : delegate_(std::move(delegate)), row_id_(row_id) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + Result NextBatchWithBitmap() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, delegate_->NextBatchWithBitmap()); + if (!IsEofBatch(batch)) { + batch.second.Add(row_id_); + } + return batch; + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + int32_t row_id_; +}; + +} // namespace + +class RealtimePrimaryKeyReaderTest : public testing::Test { + protected: + std::shared_ptr pool_ = GetDefaultPool(); +}; + +TEST_F(RealtimePrimaryKeyReaderTest, TestTransportSchemaLayout) { + arrow::FieldVector value_fields = {arrow::field("key", arrow::int64(), false), + arrow::field("value", arrow::utf8())}; + std::shared_ptr schema = MakeTransportSchema(value_fields); + + ASSERT_EQ(RealtimePrimaryKeyLayout::kValueKindIndex, 0); + ASSERT_EQ(RealtimePrimaryKeyLayout::kSequenceNumberIndex, 1); + ASSERT_EQ(RealtimePrimaryKeyLayout::kRealtimeOffsetIndex, 2); + ASSERT_EQ(RealtimePrimaryKeyLayout::kValueStartIndex, 3); + ASSERT_EQ(schema->field(0)->name(), "_VALUE_KIND"); + ASSERT_EQ(schema->field(1)->name(), "_SEQUENCE_NUMBER"); + ASSERT_EQ(schema->field(2)->name(), "_REALTIME_OFFSET"); + ASSERT_EQ(schema->field(3)->name(), "key"); + ASSERT_EQ(schema->field(4)->name(), "value"); + ASSERT_FALSE(schema->field(0)->nullable()); + ASSERT_FALSE(schema->field(1)->nullable()); + ASSERT_EQ(schema->field(2)->nullable(), SpecialFields::RealtimeOffset().Nullable()); + ASSERT_FALSE(schema->field(3)->nullable()); + ASSERT_TRUE(schema->field(4)->nullable()); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestTransportSchemaValidation) { + const std::shared_ptr valid = MakeTransportSchema({}); + std::vector invalid_fields; + + arrow::FieldVector wrong_type = valid->fields(); + wrong_type[0] = DataField::ConvertDataFieldToArrowField( + DataField(SpecialFields::ValueKind().Id(), + arrow::field("_VALUE_KIND", arrow::int32(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_type)); + + arrow::FieldVector nullable_sequence = valid->fields(); + nullable_sequence[1] = nullable_sequence[1]->WithNullable(true); + invalid_fields.push_back(std::move(nullable_sequence)); + + arrow::FieldVector wrong_offset_id = valid->fields(); + wrong_offset_id[2] = DataField::ConvertDataFieldToArrowField( + DataField(99, arrow::field("_REALTIME_OFFSET", arrow::int64(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_offset_id)); + + for (const arrow::FieldVector& fields : invalid_fields) { + ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyLayout::ValidateSchema(arrow::schema(fields)), + "transport schema field"); + } +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsCommittedPrefix) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr transport_schema = MakeTransportSchema(value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ + [0, 100, 0, 1, 10], + [0, 101, 1, 2, 20], + [0, 102, 2, 4, 40], + [0, 103, 3, 6, 60] + ])") + .ValueOrDie()); + + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(transport_array, transport_type, 2)); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(2, 4), + key_schema, value_schema, pool_)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(readers[0].get()))); + + std::vector row_kinds = {const_cast(RowKind::Insert()), + const_cast(RowKind::Insert())}; + std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; + std::vector expected = KeyValueChecker::GenerateKeyValues( + row_kinds, {102, 103}, levels, {{4}, {6}}, {{4, 40}, {6, 60}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsNegativeOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, -1, 1]])") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(1, 2), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "reader offset must be non-negative"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryOffsetCoverageAcrossReadersAndBatches) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 2, 1], [0, 11, 0, 2]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 12, 3, 3], [0, 13, 1, 4]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 4), + value_schema, value_schema, pool_)); + int64_t row_count = 0; + for (const std::unique_ptr& reader : readers) { + ASSERT_OK_AND_ASSIGN( + std::vector rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(reader.get()))); + row_count += static_cast(rows.size()); + } + ASSERT_EQ(4, row_count); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsMissingVisibleOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 2, 2]])") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(0, 3), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsDuplicateVisibleOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 11, 1, 2], [0, 12, 1, 3]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector first_rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(readers[0].get()))); + ASSERT_EQ(1, first_rows.size()); + ASSERT_NOK_WITH_MSG((ReadResultCollector::CollectKeyValueResult( + readers[1].get())), + "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsEmptyEofForVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([])").ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1), + transport_schema, OffsetRange(0, 1), value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "query readers did not cover the visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsEmptyReadersForVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_NOK_WITH_MSG( + RealtimePrimaryKeyReaderFactory::CreateForQuery(std::move(batch_readers), transport_schema, + OffsetRange(0, 1), value_schema, + value_schema, pool_), + "PK real-time store returned no query readers for a non-empty visible range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryAllowsEmptyReadersForEmptyVisibleRange) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(1, 1), + value_schema, value_schema, pool_)); + ASSERT_TRUE(readers.empty()); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryBitmapBounds) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1]])") + .ValueOrDie(); + auto batch_reader = std::make_unique( + std::make_unique(transport_array, transport_type, /*batch_size=*/1), + /*row_id=*/1); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + Result> result = + ReadResultCollector::CollectKeyValueResult(reader.get()); + ASSERT_TRUE(result.status().IsInvalid()); + ASSERT_NOK_WITH_MSG(result, + "selected row id 1 is out of bounds for realtime primary-key transport " + "batch length 1"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryRejectsPartialBitmap) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 1, 2]])") + .ValueOrDie(); + RoaringBitmap32 partial_bitmap; + partial_bitmap.Add(0); + auto batch_reader = std::make_unique( + transport_array, transport_type, partial_bitmap, /*read_batch_size=*/2); + batch_reader->EnableRandomizeBatchSize(false); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "must cover every raw transport row"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitRejectsPartialBitmap) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 0, 1], [0, 11, 1, 2]])") + .ValueOrDie(); + RoaringBitmap32 partial_bitmap; + partial_bitmap.Add(0); + auto batch_reader = std::make_unique( + transport_array, transport_type, partial_bitmap, /*read_batch_size=*/2); + batch_reader->EnableRandomizeBatchSize(false); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyCommitReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 2), + value_schema, value_schema, pool_)); + + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "must cover every raw transport row"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestQueryProjection) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr extra = MakeField("extra", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key, extra}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([[0, 10, 0, 1, 2]])") + .ValueOrDie()); + + auto query_batch_reader = + std::make_unique(transport_array, transport_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr query_reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(query_batch_reader), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector query_results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(query_reader.get()))); + ASSERT_EQ(query_results.size(), 1); + ASSERT_EQ(query_results[0].value->GetFieldCount(), 1); + ASSERT_EQ(query_results[0].value->GetInt(0), 1); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitOffsetCoverage) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr first_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 10, 2, 1], [0, 11, 0, 3]])") + .ValueOrDie(); + std::shared_ptr second_array = + arrow::ipc::internal::json::ArrayFromJSON(transport_type, + R"([[0, 12, 1, 2], [0, 13, 3, 4]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back( + std::make_unique(first_array, transport_type, /*read_batch_size=*/1)); + batch_readers.push_back( + std::make_unique(second_array, transport_type, /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 4), + value_schema, value_schema, pool_)); + int64_t row_count = 0; + for (const std::unique_ptr& reader : readers) { + ASSERT_OK_AND_ASSIGN( + std::vector rows, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(reader.get()))); + row_count += static_cast(rows.size()); + } + ASSERT_EQ(4, row_count); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestCommitRejectsEmptyReaders) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::vector> batch_readers; + + ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_), + "PK real-time store returned no commit readers for a sealed segment"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestRejectsDuplicateCommitOffset) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON( + transport_type, R"([[0, 10, 0, 1], [0, 11, 0, 2], [0, 12, 2, 3]])") + .ValueOrDie(); + std::vector> batch_readers; + batch_readers.push_back(std::make_unique(transport_array, transport_type, + /*read_batch_size=*/1)); + + ASSERT_OK_AND_ASSIGN(std::vector> readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit( + std::move(batch_readers), transport_schema, OffsetRange(0, 3), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG((ReadResultCollector::CollectKeyValueResult( + readers[0].get())), + "did not cover the sealed range"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestBadCommitBatch) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value = MakeField("value", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key, value}); + std::shared_ptr transport_schema = MakeTransportSchema({key, value}); + std::shared_ptr actual_schema = MakeTransportSchema({key}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyCommitReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + arrow::schema({key}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "field count"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestSafeDecode) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr transport_schema = MakeTransportSchema({key}); + + arrow::FieldVector invalid_fields = transport_schema->fields(); + invalid_fields[0] = invalid_fields[0]->WithName("wrong_value_kind"); + invalid_fields[3] = MakeField("wrong_key", arrow::int32(), 99); + std::shared_ptr invalid_type = arrow::struct_(invalid_fields); + auto invalid_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[0, 10, 0, 1]])").ValueOrDie()); + + auto batch_reader = std::make_unique(invalid_array, invalid_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "transport batch field"); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestNestedValues) { + std::shared_ptr id = MakeField("id", arrow::int32(), 0); + std::shared_ptr key_schema = arrow::schema({id}); + std::shared_ptr query_item_b = MakeField("renamed_b", arrow::int32(), 11); + std::shared_ptr query_item_a = MakeField("renamed_a", arrow::int32(), 10); + std::shared_ptr query_items = MakeField( + "items_renamed", + arrow::list(arrow::field("element", arrow::struct_({query_item_b, query_item_a}))), 2); + std::shared_ptr query_attr_y = MakeField("renamed_y", arrow::int32(), 21); + std::shared_ptr query_attr_x = MakeField("renamed_x", arrow::int32(), 20); + std::shared_ptr query_attrs = + MakeField("attrs_renamed", + arrow::map(arrow::utf8(), arrow::struct_({query_attr_y, query_attr_x})), 3); + std::shared_ptr query_key_right = MakeField("renamed_right", arrow::int32(), 31); + std::shared_ptr query_key_left = MakeField("renamed_left", arrow::int32(), 30); + std::shared_ptr query_keyed_values = + MakeField("keyed_values_renamed", + arrow::map(arrow::struct_({query_key_right, query_key_left}), arrow::int32()), 4); + std::shared_ptr query_value_schema = + arrow::schema({id, query_items, query_attrs, query_keyed_values}); + std::shared_ptr transport_schema = + MakeTransportSchema(query_value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + std::shared_ptr transport_array = + arrow::ipc::internal::json::ArrayFromJSON( + transport_type, + R"([[0, 10, 0, 1, [[200, 100], [400, 300]], [["k1", [8, 7]], ["k2", [10, 9]]], [[[12, 11], 13], [[22, 21], 23]]]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(transport_array, transport_type, 1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + CreateRealtimePrimaryKeyQueryReaderForTest( + std::move(batch_reader), transport_schema, OffsetRange(0, 1), + key_schema, query_value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0].key->GetInt(0), 1); + ASSERT_EQ(results[0].value->GetFieldCount(), 4); + ASSERT_EQ(results[0].value->GetInt(0), 1); + + std::shared_ptr item_array = results[0].value->GetArray(1); + ASSERT_EQ(item_array->Size(), 2); + std::shared_ptr first_item = item_array->GetRow(0, 2); + ASSERT_EQ(first_item->GetInt(0), 200); + ASSERT_EQ(first_item->GetInt(1), 100); + std::shared_ptr second_item = item_array->GetRow(1, 2); + ASSERT_EQ(second_item->GetInt(0), 400); + ASSERT_EQ(second_item->GetInt(1), 300); + + std::shared_ptr attr_map = results[0].value->GetMap(2); + ASSERT_EQ(attr_map->Size(), 2); + std::shared_ptr key_array = attr_map->KeyArray(); + ASSERT_EQ(std::string(key_array->GetStringView(0)), "k1"); + ASSERT_EQ(std::string(key_array->GetStringView(1)), "k2"); + std::shared_ptr value_array = attr_map->ValueArray(); + std::shared_ptr first_attr = value_array->GetRow(0, 2); + ASSERT_EQ(first_attr->GetInt(0), 8); + ASSERT_EQ(first_attr->GetInt(1), 7); + std::shared_ptr second_attr = value_array->GetRow(1, 2); + ASSERT_EQ(second_attr->GetInt(0), 10); + ASSERT_EQ(second_attr->GetInt(1), 9); + + std::shared_ptr keyed_value_map = results[0].value->GetMap(3); + ASSERT_EQ(keyed_value_map->Size(), 2); + std::shared_ptr struct_keys = keyed_value_map->KeyArray(); + std::shared_ptr first_key = struct_keys->GetRow(0, 2); + ASSERT_EQ(first_key->GetInt(0), 12); + ASSERT_EQ(first_key->GetInt(1), 11); + std::shared_ptr second_key = struct_keys->GetRow(1, 2); + ASSERT_EQ(second_key->GetInt(0), 22); + ASSERT_EQ(second_key->GetInt(1), 21); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(0), 13); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(1), 23); +} + +TEST_F(RealtimePrimaryKeyReaderTest, TestFactoryFailureClosesReaders) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr transport_schema = MakeTransportSchema(value_schema->fields()); + std::shared_ptr transport_type = arrow::struct_(transport_schema->fields()); + auto transport_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(transport_type, R"([ + [0, 10, 0, 1, 100] + ])") + .ValueOrDie()); + + int32_t factory_failure_close_count = 0; + std::vector> batch_readers; + batch_readers.push_back(std::make_unique( + std::make_unique(transport_array, transport_type, 1), + &factory_failure_close_count)); + batch_readers.push_back(nullptr); + ASSERT_NOK_WITH_MSG(RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_), + "PK real-time store returned a null query reader"); + ASSERT_EQ(factory_failure_close_count, 1); +} + +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp new file mode 100644 index 000000000..61cd14009 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,293 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/realtime_primary_key_writer.h" + +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/core_options.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/core/utils/primary_key_table_utils.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +Result> CreateRealtimePrimaryKeyTransportBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, int64_t first_sequence_number, + int64_t first_offset, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr input, + arrow::ImportArray(batch->GetData(), arrow::struct_(write_schema->fields()))); + if (!input || input->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr values = checked_pointer_cast(input); + const int64_t count = values->length(); + arrow::Int8Builder kinds(arrow_pool); + arrow::Int64Builder sequences(arrow_pool); + arrow::Int64Builder offsets(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count)); + const std::vector& row_kinds = batch->GetRowKind(); + for (int64_t row = 0; row < count; ++row) { + const RecordBatch::RowKind kind = + row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row]; + kinds.UnsafeAppend(static_cast(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr kind_array; + std::shared_ptr sequence_array; + std::shared_ptr offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array)); + arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array), + std::move(offset_array)}; + columns.insert(columns.end(), values->fields().begin(), values->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr transport, + arrow::StructArray::Make(std::move(columns), transport_schema->fields())); + + std::vector sort_keys; + sort_keys.reserve(trimmed_primary_keys.size() + 1); + for (const std::string& key : trimmed_primary_keys) { + sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + } + sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), + arrow::compute::SortOrder::Ascending); + arrow::compute::ExecContext context(arrow_pool); + arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum indices, + arrow::compute::SortIndices(arrow::Datum(transport), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(transport), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + return checked_pointer_cast(sorted.make_array()); +} + +} // namespace + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool) { + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr field = write_schema->GetFieldByName(key); + if (!field) { + return Status::Invalid("PK field is missing from write schema: ", key); + } + key_fields.push_back(std::move(field)); + } + const RealtimePartitionBucket partition_bucket(partition, bucket); + PAIMON_ASSIGN_OR_RAISE(int64_t initial_max_sequence_number, + realtime_context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, restored_max_sequence_number)); + return std::shared_ptr(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, + transport_schema, arrow::schema(std::move(key_fields)), trimmed_primary_keys, + key_comparator, options, store_state.initial_offset, initial_max_sequence_number, + memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + int64_t next_offset, int64_t last_sequence_number, + const std::shared_ptr& memory_pool) + : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + realtime_store_(realtime_store), + merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), + write_schema_(write_schema), + transport_schema_(transport_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + options_(options), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t count = batch->GetData()->length; + if (count == 0) { + return Status::OK(); + } + const std::vector& row_kinds = batch->GetRowKind(); + if (!row_kinds.empty() && static_cast(row_kinds.size()) != count) { + return Status::Invalid("PK real-time row-kind count does not match batch row count"); + } + for (RecordBatch::RowKind row_kind : row_kinds) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, + RowKind::FromByteValue(static_cast(row_kind))); + static_cast(validated); + } + std::lock_guard lock(realtime_store_mutex_); + if (count > std::numeric_limits::max() - next_offset_) { + return Status::Invalid("real-time offset range exceeds INT64_MAX"); + } + // Reserve INT64_MAX as the exhausted sequence-number sentinel. + if (last_sequence_number_ >= std::numeric_limits::max() - count) { + return Status::Invalid("PK sequence range exceeds INT64_MAX"); + } + const int64_t first_sequence = last_sequence_number_ + 1; + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr transport, + CreateRealtimePrimaryKeyTransportBatch(std::move(batch), write_schema_, transport_schema_, + trimmed_primary_keys_, first_sequence, next_offset_, + arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport, output.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(transport_batch), OffsetRange(next_offset_, next_offset_ + count)})); + next_offset_ += count; + last_sequence_number_ += count; + PAIMON_RETURN_NOT_OK(realtime_context_->AdvanceMaterializedMaxSequenceNumber( + partition_bucket_, last_sequence_number_)); + return Status::OK(); +} + +Result RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard prepare_lock(prepare_mutex_); + std::optional> segment; + { + std::lock_guard store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> sealed, + realtime_store_->SealForCommit()); + segment = std::move(sealed); + } + if (segment && !segment.value()) { + return Status::Invalid("PK real-time store sealed a null segment"); + } + std::optional sealed_range; + if (segment) { + sealed_range = segment.value()->GetOffsetRange(); + if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), sealed_range.value())); + } + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, + merge_tree_writer_->PrepareCommit(wait_compaction)); + if (segment) { + increment.SetRealtimeOffsetRange(sealed_range.value()); + } + return increment; +} + +Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr& segment, + const OffsetRange& sealed_offsets) { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + realtime_store_->CreateCommitReaders(segment)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> realtime_primary_key_readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit(std::move(readers), transport_schema_, + sealed_offsets, key_schema_, write_schema_, + memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(realtime_primary_key_readers.size()); + for (std::unique_ptr& realtime_primary_key_reader : + realtime_primary_key_readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + write_schema_, trimmed_primary_keys_, options_, memory_pool_)); + sorted_readers.push_back(std::make_unique( + std::move(realtime_primary_key_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); + } + return merge_tree_writer_->WriteSortedReadersToFiles(std::move(sorted_readers)); +} + +Status RealtimePrimaryKeyWriter::Compact(bool) { + return Status::Invalid("PK real-time write does not support explicit compaction"); +} +uint64_t RealtimePrimaryKeyWriter::GetMemoryUsage() const { + return realtime_store_->GetMemoryUsage(); +} +Status RealtimePrimaryKeyWriter::FlushMemory() { + return Status::OK(); +} +Result RealtimePrimaryKeyWriter::CompactNotCompleted() { + return merge_tree_writer_->CompactNotCompleted(); +} +Status RealtimePrimaryKeyWriter::Sync() { + return merge_tree_writer_->Sync(); +} +Status RealtimePrimaryKeyWriter::Close() { + return merge_tree_writer_->Close(); +} +std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { + return merge_tree_writer_->GetMetrics(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h new file mode 100644 index 000000000..cdd3d889f --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "paimon/core/core_options.h" +#include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" + +namespace arrow { +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; +class MergeTreeWriter; +class FieldsComparator; +class RealtimeContextImpl; +struct RealtimeStoreState; + +class RealtimePrimaryKeyWriter final : public BatchWriter { + public: + static Result> Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool); + + Status Write(std::unique_ptr&& batch) override; + Result PrepareCommit(bool wait_compaction) override; + Status Compact(bool full_compaction) override; + uint64_t GetMemoryUsage() const override; + Status FlushMemory() override; + Result CompactNotCompleted() override; + Status Sync() override; + Status Close() override; + std::shared_ptr GetMetrics() const override; + + private: + RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const CoreOptions& options, int64_t next_offset, + int64_t last_sequence_number, + const std::shared_ptr& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment, + const OffsetRange& sealed_offsets); + + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr realtime_store_; + std::shared_ptr merge_tree_writer_; + std::shared_ptr realtime_context_; + RealtimePartitionBucket partition_bucket_; + std::shared_ptr write_schema_; + std::shared_ptr transport_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; + CoreOptions options_; + int64_t next_offset_; + int64_t last_sequence_number_; + std::mutex realtime_store_mutex_; + std::mutex prepare_mutex_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_reader_test.cpp b/src/paimon/core/realtime/realtime_reader_test.cpp index ec37cfed4..ded060989 100644 --- a/src/paimon/core/realtime/realtime_reader_test.cpp +++ b/src/paimon/core/realtime/realtime_reader_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include "paimon/arrow/abi.h" #include "paimon/testing/utils/testharness.h" @@ -37,6 +38,8 @@ class TestingReadView : public RealtimeReadView { class TestingBatchReader : public BatchReader { public: + explicit TestingBatchReader(int32_t* close_count = nullptr) : close_count_(close_count) {} + Result NextBatch() override { return MakeEofBatch(); } @@ -45,7 +48,14 @@ class TestingBatchReader : public BatchReader { return nullptr; } - void Close() override {} + void Close() override { + if (close_count_) { + ++(*close_count_); + } + } + + private: + int32_t* close_count_; }; TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { @@ -57,5 +67,19 @@ TEST(RealtimeReaderTest, TestRejectsIncompleteReader) { "inner reader is null"); } +TEST(RealtimeReaderTest, TestCloseReleasesResources) { + int32_t close_count = 0; + std::shared_ptr read_view = std::make_shared(); + std::weak_ptr weak_read_view = read_view; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + RealtimeReader::Create(std::move(read_view), + std::make_unique(&close_count))); + ASSERT_FALSE(weak_read_view.expired()); + reader->Close(); + ASSERT_EQ(1, close_count); + ASSERT_TRUE(weak_read_view.expired()); +} + } // namespace } // namespace paimon::test diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 578856260..4599fc325 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -46,6 +46,14 @@ TEST(SchemaValidationTest, TestSimple) { ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestRealtimeOffsetIsGloballyReserved) { + auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(0, schema, {}, {}, {})); + ASSERT_NOK_WITH_MSG(SchemaValidation::ValidateTableSchema(*table_schema), + "field name '_REALTIME_OFFSET' in schema cannot be special field"); +} + TEST(SchemaValidationTest, TestVectorType) { auto vector_field = arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)); auto schema = arrow::schema({arrow::field("id", arrow::int64()), vector_field}); diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index 6885dc374..12b3823a7 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -77,6 +77,13 @@ Result> AppendOnlyTableRead::CreateReader( std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::shared_ptr& split : splits) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); @@ -92,8 +99,6 @@ Result> AppendOnlyTableRead::CreateReader( } } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (!realtime_splits.empty()) { const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { @@ -106,7 +111,7 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } - return result; + return std::make_unique(std::move(readers), GetMemoryPool()); } Result> AppendOnlyTableRead::CreateRealtimeReader( @@ -124,6 +129,13 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); std::vector> readers; readers.reserve(realtime_split->DiskSplits().size() + 1); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), realtime_split->Bucket()); if (memory.partition_bucket != expected_partition_bucket) { @@ -150,8 +162,17 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( std::vector> memory_readers, memory.store->CreateQueryReaders(memory.read_view, realtime_split->CommittedEndOffset(), query_context)); - + const size_t first_memory_reader = readers.size(); + readers.reserve(readers.size() + memory_readers.size()); for (std::unique_ptr& memory_reader : memory_readers) { + readers.push_back(std::move(memory_reader)); + } + + for (size_t i = first_memory_reader; i < readers.size(); ++i) { + std::unique_ptr& memory_reader = readers[i]; + if (!memory_reader) { + return Status::Invalid("append-only real-time store returned a null query reader"); + } if (context_->EnablePredicateFilter() && context_->GetPredicate()) { PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create( std::move(memory_reader), @@ -159,14 +180,15 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, RealtimeReader::Create(memory.read_view, std::move(memory_reader))); - readers.push_back(std::move(realtime_reader)); + memory_reader = std::move(realtime_reader); } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (release_ticket) { PAIMON_RETURN_NOT_OK( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + readers_guard.Release(); return result; } diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 208807493..afb852a6b 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -19,13 +19,31 @@ #include "paimon/core/table/source/key_value_table_read.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/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/key_value.h" +#include "paimon/core/mergetree/compact/merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/realtime/realtime_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" +#include "paimon/core/table/source/realtime_split.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/status.h" namespace paimon { @@ -34,16 +52,79 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +struct ColumnarBatchContext; -KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, - const std::shared_ptr& path_factory, - const std::shared_ptr& context, - const std::shared_ptr& memory_pool, - const std::shared_ptr& executor) +namespace { + +Result> CreateRealtimePrimaryKeyQueryTransportSchema( + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema) { + arrow::FieldVector transport_value_fields; + transport_value_fields.reserve(key_schema->num_fields() + value_schema->num_fields()); + std::unordered_set field_ids; + for (const std::shared_ptr& field : key_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field_ids.insert(field_id).second) { + transport_value_fields.push_back(field); + } + } + for (const std::shared_ptr& field : value_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field_ids.insert(field_id).second) { + transport_value_fields.push_back(field); + } + } + return RealtimePrimaryKeyLayout::CreateSchema(transport_value_fields); +} + +Result>> CreateMemoryReaders( + const std::shared_ptr& split, const RealtimePartitionBucketView& memory, + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& context, + const std::shared_ptr& memory_pool) { + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*transport_schema, c_schema.get())); + ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, + memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> realtime_primary_key_readers, + RealtimePrimaryKeyReaderFactory::CreateForQuery( + std::move(batch_readers), transport_schema, + OffsetRange(split->CommittedEndOffset(), split->MemoryEndOffset()), key_schema, + value_schema, memory_pool)); + std::vector> result; + result.reserve(realtime_primary_key_readers.size()); + for (std::unique_ptr& realtime_primary_key_reader : + realtime_primary_key_readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge, + PrimaryKeyTableUtils::CreateMergeFunction( + value_schema, context->GetTableSchema()->PrimaryKeys(), + context->GetCoreOptions(), memory_pool)); + result.push_back(std::make_unique( + std::move(realtime_primary_key_reader), key_comparator, + std::make_shared(std::move(merge)))); + } + return result; +} + +} // namespace + +KeyValueTableRead::KeyValueTableRead( + std::vector>&& split_reads, + const std::shared_ptr& path_factory, + const std::shared_ptr& context, + const std::shared_ptr& realtime_primary_key_transport_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& executor) : TableRead(memory_pool), split_reads_(std::move(split_reads)), path_factory_(path_factory), context_(context), + realtime_primary_key_transport_schema_(realtime_primary_key_transport_schema), executor_(executor) {} Result> KeyValueTableRead::Create( @@ -57,10 +138,18 @@ Result> KeyValueTableRead::Create( PAIMON_ASSIGN_OR_RAISE( std::unique_ptr merge_file_split_read, MergeFileSplitRead::Create(path_factory, context, memory_pool, executor)); + std::shared_ptr realtime_primary_key_transport_schema; + if (context->GetRealtimeContext()) { + PAIMON_ASSIGN_OR_RAISE( + realtime_primary_key_transport_schema, + CreateRealtimePrimaryKeyQueryTransportSchema(merge_file_split_read->GetKeySchema(), + merge_file_split_read->GetValueSchema())); + } split_reads.emplace_back(std::move(merge_file_split_read)); - return std::unique_ptr(new KeyValueTableRead(std::move(split_reads), path_factory, - context, memory_pool, executor)); + return std::unique_ptr( + new KeyValueTableRead(std::move(split_reads), path_factory, context, + realtime_primary_key_transport_schema, memory_pool, executor)); } void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { @@ -75,6 +164,11 @@ void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { + std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); + if (realtime_split) { + return CreateRealtimeReader(realtime_split, /*release_ticket=*/true); + } + std::shared_ptr dispatch_split = split; if (auto indexed_split = std::dynamic_pointer_cast(split)) { PAIMON_RETURN_NOT_OK(indexed_split->Validate()); @@ -126,8 +220,104 @@ Result> KeyValueTableRead::CreateReader( return Status::Invalid("create reader failed, not read match with data split."); } +Result> KeyValueTableRead::CreateReader( + const std::vector>& splits) { + std::vector> readers; + readers.reserve(splits.size()); + std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); + for (const std::shared_ptr& split : splits) { + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(split); + if (realtime_split) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + CreateRealtimeReader(realtime_split, + /*release_ticket=*/false)); + readers.push_back(std::move(reader)); + realtime_splits.push_back(std::move(realtime_split)); + } else { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, CreateReader(split)); + readers.push_back(std::move(reader)); + } + } + + if (!realtime_splits.empty()) { + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + for (const std::shared_ptr& realtime_split : realtime_splits) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + } + return std::make_unique(std::move(readers), GetMemoryPool()); +} + +Result> KeyValueTableRead::CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket) { + if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { + return Status::Invalid("unsupported real-time split version"); + } + if (realtime_split->MemoryEndOffset() < realtime_split->CommittedEndOffset()) { + return Status::Invalid("real-time split memory end offset precedes committed end offset"); + } + const std::shared_ptr realtime_context = context_->GetRealtimeContext(); + if (!realtime_context) { + return Status::Invalid("reading a real-time split requires a real-time context"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(RealtimePartitionBucketView memory, + realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); + const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), + realtime_split->Bucket()); + if (memory.partition_bucket != expected_partition_bucket) { + return Status::Invalid("real-time read-view ticket belongs to another partition-bucket"); + } + const std::optional memory_range = memory.read_view->GetOffsetRange(); + if (!memory_range || memory_range->end != realtime_split->MemoryEndOffset()) { + return Status::Invalid("real-time read-view ticket does not match the split offset range"); + } + for (const std::unique_ptr& read : split_reads_) { + auto* merge_read = dynamic_cast(read.get()); + if (merge_read) { + PAIMON_ASSIGN_OR_RAISE( + std::vector> memory_readers, + CreateMemoryReaders(realtime_split, memory, realtime_primary_key_transport_schema_, + merge_read->GetKeySchema(), merge_read->GetValueSchema(), + merge_read->GetKeyComparator(), context_, GetMemoryPool())); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + merge_read->CreateRealtimeReader(realtime_split->DiskSplits(), + std::move(memory_readers))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, + RealtimeReader::Create(memory.read_view, std::move(reader))); + if (release_ticket) { + PAIMON_RETURN_NOT_OK( + realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); + } + return std::unique_ptr(std::move(realtime_reader)); + } + } + return Status::Invalid("create reader failed, merge file split read not found"); +} + Result> KeyValueTableRead::CreateCountReader( const std::vector>& splits) { + for (const std::shared_ptr& split : splits) { + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented( + "CreateCountReader does not support process-local real-time splits"); + } + } if (context_->GetPredicate() != nullptr) { return Status::NotImplemented( "CreateCountReader with predicate pushdown is not supported yet"); diff --git a/src/paimon/core/table/source/key_value_table_read.h b/src/paimon/core/table/source/key_value_table_read.h index d6a1c83d3..1dd59b016 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -22,6 +22,7 @@ #include #include +#include "arrow/type_fwd.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/operation/split_read.h" #include "paimon/core/utils/file_store_path_factory.h" @@ -35,6 +36,7 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +class RealtimeSplit; class KeyValueTableRead : public TableRead { public: @@ -45,6 +47,9 @@ class KeyValueTableRead : public TableRead { Result> CreateReader(const std::shared_ptr& split) override; + Result> CreateReader( + const std::vector>& splits) override; + Result> CreateCountReader( const std::vector>& splits) override; @@ -54,12 +59,17 @@ class KeyValueTableRead : public TableRead { KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, const std::shared_ptr& context, + const std::shared_ptr& realtime_primary_key_transport_schema, const std::shared_ptr& memory_pool, const std::shared_ptr& executor); + Result> CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket); + std::vector> split_reads_; std::shared_ptr path_factory_; std::shared_ptr context_; + std::shared_ptr realtime_primary_key_transport_schema_; std::shared_ptr executor_; bool force_keep_delete_ = false; }; diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index c275208c5..4c3968dc3 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -108,7 +108,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl for (const std::shared_ptr& split : disk_splits) { std::shared_ptr data_split = std::dynamic_pointer_cast(split); if (!data_split) { - return Status::Invalid("real-time append scan requires process-local data splits"); + return Status::Invalid("real-time scan requires process-local data splits"); } std::vector> partition_values; PAIMON_ASSIGN_OR_RAISE(partition_values, @@ -152,16 +152,17 @@ Result>> RealtimeTableScan::CreateRealtimeSpl continue; } - // Append tables can schedule all but the tail disk split independently. The tail split - // carries the immutable memory view so disk and memory are still concatenated by one - // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit. - auto tail_disk_split = std::prev(grouped_disk_splits.end()); - result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split); - std::vector> realtime_disk_splits; - realtime_disk_splits.push_back(std::move(*tail_disk_split)); RealtimePartitionBucketView& memory = memory_iter->second; + if (!pk_table_) { + // Append tables can schedule all but the tail disk split independently. The tail split + // carries the immutable memory view so disk and memory are still concatenated by one + // RealtimeSplit without collapsing the whole partition-bucket into one scheduling unit. + auto tail_disk_split = std::prev(grouped_disk_splits.end()); + result.insert(result.end(), grouped_disk_splits.begin(), tail_disk_split); + grouped_disk_splits.erase(grouped_disk_splits.begin(), tail_disk_split); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_split, - create_realtime_split(key, std::move(realtime_disk_splits), memory)); + create_realtime_split(key, std::move(grouped_disk_splits), memory)); result.push_back(std::move(realtime_split)); active_memory.erase(memory_iter); } @@ -176,7 +177,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl return result; } -RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, +RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -184,6 +185,7 @@ RealtimeTableScan::RealtimeTableScan(std::unique_ptr&& disk_scan, const std::shared_ptr& scan_filter, int64_t read_view_ttl_millis) : disk_scan_(std::move(disk_scan)), + pk_table_(pk_table), realtime_context_(realtime_context), path_factory_(path_factory), snapshot_manager_(snapshot_manager), diff --git a/src/paimon/core/table/source/realtime_table_scan.h b/src/paimon/core/table/source/realtime_table_scan.h index 7d036d420..692b749ef 100644 --- a/src/paimon/core/table/source/realtime_table_scan.h +++ b/src/paimon/core/table/source/realtime_table_scan.h @@ -35,10 +35,10 @@ class FileSystem; class ScanFilter; class SnapshotManager; -/// Adds process-local memory splits to a normal append-table batch scan. +/// Adds process-local memory splits to a normal data-table batch scan. class RealtimeTableScan : public TableScan { public: - RealtimeTableScan(std::unique_ptr&& disk_scan, + RealtimeTableScan(std::unique_ptr&& disk_scan, bool pk_table, const std::shared_ptr& realtime_context, const std::shared_ptr& path_factory, const std::shared_ptr& snapshot_manager, @@ -67,6 +67,7 @@ class RealtimeTableScan : public TableScan { const std::optional& snapshot_id) const; std::unique_ptr disk_scan_; + bool pk_table_; std::shared_ptr realtime_context_; std::shared_ptr path_factory_; std::shared_ptr snapshot_manager_; diff --git a/src/paimon/core/table/source/table_scan.cpp b/src/paimon/core/table/source/table_scan.cpp index 2dda955ac..66a3f7426 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -63,6 +63,7 @@ #include "paimon/core/utils/field_mapping.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/index_file_path_factories.h" +#include "paimon/core/utils/primary_key_table_utils.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/format/file_format.h" #include "paimon/realtime/realtime_context.h" @@ -218,22 +219,27 @@ Result> TableScan::Create(std::unique_ptr> NewDataTableScan(const std::shared_ptrGetSpecificFileSystem(), {})); core_options.WithCache(context->GetCache()); - PAIMON_RETURN_NOT_OK(ValidateRealtimeScan(*table_schema, core_options, *context)); + PAIMON_RETURN_NOT_OK( + ValidateRealtimeScan(*table_schema, core_options, *context, read_optimized)); // validate options if (core_options.GetBucket() == -1) { if (!table_schema->PrimaryKeys().empty()) { @@ -343,7 +350,7 @@ Result> NewDataTableScan(const std::shared_ptr realtime_context, RealtimeContextImpl::Cast(context->GetRealtimeContext())); return std::make_unique( - std::move(batch_scan), realtime_context, path_factory, + std::move(batch_scan), pk_table, realtime_context, path_factory, snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(), context->GetScanFilters(), core_options.GetRealtimeReadViewTtlMillis()); } diff --git a/src/paimon/core/table/system/read_optimized_system_table.cpp b/src/paimon/core/table/system/read_optimized_system_table.cpp index 6abec946b..d7bfa0912 100644 --- a/src/paimon/core/table/system/read_optimized_system_table.cpp +++ b/src/paimon/core/table/system/read_optimized_system_table.cpp @@ -58,6 +58,10 @@ std::map ReadOptimizedSystemTable::ReadOptimizedOption Result> ReadOptimizedSystemTable::NewScan( const std::shared_ptr& context) const { + if (context->GetRealtimeContext() && !table_schema_->PrimaryKeys().empty()) { + return Status::NotImplemented( + "PK real-time union read does not support read-optimized scans"); + } auto options = ReadOptimizedOptions(); ScanContextBuilder builder(table_path_); builder.SetOptions(options) diff --git a/src/paimon/core/utils/primary_key_table_utils.cpp b/src/paimon/core/utils/primary_key_table_utils.cpp index cf72da4ae..2ca444a7a 100644 --- a/src/paimon/core/utils/primary_key_table_utils.cpp +++ b/src/paimon/core/utils/primary_key_table_utils.cpp @@ -29,12 +29,14 @@ #include "paimon/common/utils/fields_comparator.h" #include "paimon/common/utils/object_utils.h" #include "paimon/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" #include "paimon/core/mergetree/compact/aggregate/aggregate_merge_function.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/first_row_merge_function.h" #include "paimon/core/mergetree/compact/merge_function.h" #include "paimon/core/mergetree/compact/partial_update_merge_function.h" #include "paimon/core/options/merge_engine.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/status.h" namespace arrow { @@ -96,4 +98,52 @@ Result> PrimaryKeyTableUtils::CreateSequenceFi options.SequenceFieldSortOrderIsAscending()); } +Status PrimaryKeyTableUtils::ValidateRealtimeOptions(const CoreOptions& options, + const TableSchema& schema) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime does not support data evolution"); + } + if (options.IgnoreDelete()) { + return Status::NotImplemented("PK realtime requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime supports only ascending sequence.field.sort-order"); + } + if (options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime supports only the NONE changelog producer"); + } + if (options.DeletionVectorsEnabled()) { + return Status::NotImplemented("PK realtime does not support deletion vectors"); + } + if (options.NeedLookup()) { + return Status::NotImplemented("PK realtime does not support lookup"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector primary_key_fields, + schema.TrimmedPrimaryKeyFields()); + for (const DataField& field : primary_key_fields) { + if (field.Type()->id() == arrow::Type::FLOAT || field.Type()->id() == arrow::Type::DOUBLE) { + return Status::NotImplemented( + "PK realtime does not support FLOAT or DOUBLE primary keys"); + } + } + if (options.GlobalIndexEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(schema)); + if (!definitions.Definitions().empty()) { + return Status::NotImplemented("PK realtime does not support global indexes"); + } + } + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils.h b/src/paimon/core/utils/primary_key_table_utils.h index 82a108ab7..114801cc1 100644 --- a/src/paimon/core/utils/primary_key_table_utils.h +++ b/src/paimon/core/utils/primary_key_table_utils.h @@ -36,6 +36,7 @@ class CoreOptions; class MemoryPool; class FieldsComparator; class DataField; +class TableSchema; class PrimaryKeyTableUtils { public: @@ -57,6 +58,8 @@ class PrimaryKeyTableUtils { static Result> CreateSequenceFieldsComparator( const std::vector& value_fields, const CoreOptions& options); + + static Status ValidateRealtimeOptions(const CoreOptions& options, const TableSchema& schema); }; } // namespace paimon diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 12713ca5b..796922336 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -22,7 +22,9 @@ #include #include #include +#include #include +#include #include "arrow/type.h" #include "gtest/gtest.h" @@ -32,6 +34,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/key_value.h" #include "paimon/core/mergetree/compact/merge_function.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/defs.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" @@ -40,6 +43,110 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +std::shared_ptr PkSchema( + const std::shared_ptr& key_type = arrow::int64(), + const std::map& options = {}) { + return TableSchema::Create( + /*schema_id=*/0, + arrow::schema({arrow::field("id", key_type), arrow::field("value", arrow::utf8())}), + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options) + .value(); +} + +} // namespace + +TEST(PrimaryKeyTableUtilsTest, TestSupportedRealtimeOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema())); +} + +TEST(PrimaryKeyTableUtilsTest, TestUnsupportedRealtimeOptions) { + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema())); + } +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeAcceptsInactiveMergeEngineOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::map option_map = { + {Options::BUCKET, "1"}, + {sequence_group, "seq"}, + {Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_DELETE, "true"}, + {Options::AGGREGATION_REMOVE_RECORD_ON_DELETE, "true"}, + {Options::PARTIAL_UPDATE_REMOVE_RECORD_ON_SEQUENCE_GROUP, "seq"}, + }; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr schema, + TableSchema::Create(0, + arrow::schema({arrow::field("id", arrow::int64()), + arrow::field("value", arrow::utf8()), + arrow::field("seq", arrow::int64())}), + {}, {"id"}, option_map)); + ASSERT_OK(PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *schema)); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsDeleteAndSequenceOrderingOptions) { + ASSERT_OK_AND_ASSIGN( + CoreOptions ignore_delete, + CoreOptions::FromMap({{Options::BUCKET, "1"}, {Options::IGNORE_DELETE, "true"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyTableUtils::ValidateRealtimeOptions(ignore_delete, *PkSchema()), + "requires default delete behavior"); + + ASSERT_OK_AND_ASSIGN( + CoreOptions descending, + CoreOptions::FromMap( + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD_SORT_ORDER, "descending"}})); + ASSERT_NOK_WITH_MSG(PrimaryKeyTableUtils::ValidateRealtimeOptions(descending, *PkSchema()), + "supports only ascending sequence.field.sort-order"); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeReportsSpecificLookupErrors) { + const std::vector, std::string>> cases = { + {{{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + "PK realtime does not support lookup"}, + {{{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + "PK realtime does not support deletion vectors"}, + {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + "PK realtime supports only the NONE changelog producer"}, + {{{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "lookup"}}, + "PK realtime supports only the NONE changelog producer"}, + }; + for (const auto& [option_map, expected_message] : cases) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + Status status = PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema()); + ASSERT_TRUE(status.IsNotImplemented()); + ASSERT_EQ(status.message(), expected_message); + } +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsFloatingPrimaryKeys) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_NOK_WITH_MSG( + PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema(arrow::float32())), + "FLOAT or DOUBLE primary keys"); + ASSERT_NOK_WITH_MSG( + PrimaryKeyTableUtils::ValidateRealtimeOptions(options, *PkSchema(arrow::float64())), + "FLOAT or DOUBLE primary keys"); +} + +TEST(PrimaryKeyTableUtilsTest, TestRealtimeRejectsEnabledGlobalIndex) { + const std::map option_map = {{Options::BUCKET, "1"}, + {Options::PK_BTREE_INDEX_COLUMNS, "id"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK_WITH_MSG(PrimaryKeyTableUtils::ValidateRealtimeOptions( + options, *PkSchema(arrow::int64(), option_map)), + "does not support global indexes"); +} TEST(PrimaryKeyTableUtilsTest, TestCreateSequenceFieldsComparator) { { diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 6298137ea..4e28f286f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -42,16 +43,23 @@ #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/commit_context.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/table_schema.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/core/table/source/realtime_split.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/fs/file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/orphan_files_cleaner.h" #include "paimon/predicate/function.h" @@ -59,6 +67,7 @@ #include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" #include "paimon/reader/count_reader.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_store.h" #include "paimon/record_batch.h" @@ -71,6 +80,129 @@ #include "paimon/write_context.h" namespace paimon::test { +namespace { + +class TrackingRealtimeReadView final : public RealtimeReadView { + public: + explicit TrackingRealtimeReadView(std::shared_ptr delegate) + : delegate_(std::move(delegate)) {} + + std::optional GetOffsetRange() const override { + return delegate_->GetOffsetRange(); + } + + const std::shared_ptr& Delegate() const { + return delegate_; + } + + private: + std::shared_ptr delegate_; +}; + +class DelegatingRealtimeStore : public RealtimeStore { + public: + explicit DelegatingRealtimeStore(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + protected: + std::shared_ptr delegate_; +}; + +class DecoratingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + using Decorator = + std::function(const std::shared_ptr&)>; + + explicit DecoratingRealtimeStoreFactory(Decorator decorator) + : decorator_(std::move(decorator)) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return decorator_(delegate); + } + + private: + ArrowRealtimeStoreFactory delegate_; + Decorator decorator_; +}; + +template +std::shared_ptr MakeDecoratingFactory(Args... args) { + return std::make_shared( + [=](const std::shared_ptr& delegate) -> std::shared_ptr { + return std::make_shared(delegate, args...); + }); +} + +class QueryTrackingRealtimeStore final : public DelegatingRealtimeStore { + public: + QueryTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : DelegatingRealtimeStore(delegate), + saw_query_predicate_(saw_query_predicate), + query_view_(query_view) {} + + Result> AcquireReadView() override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate_view, + delegate_->AcquireReadView()); + return std::shared_ptr( + std::make_shared(delegate_view)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + if (context.predicate) { + saw_query_predicate_->store(true, std::memory_order_release); + } + *query_view_ = view; + std::shared_ptr tracking_view = + std::dynamic_pointer_cast(view); + if (!tracking_view) { + return Status::Invalid("query tracking store received an unexpected read view"); + } + return delegate_->CreateQueryReaders(tracking_view->Delegate(), offset_begin, context); + } + + private: + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +} // namespace namespace { @@ -219,6 +351,20 @@ class RealtimeWriteInteTest : public ::testing::Test { /*ignore_if_exists=*/false)); } + void CreatePkTable(const std::vector& partition_keys = {}, + const std::vector& primary_keys = {"id"}) const { + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema_, c_schema.get()).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr catalog, + Catalog::Create(dir_->Str(), options_)); + ASSERT_OK(catalog->CreateDatabase("foo", {}, /*ignore_if_exists=*/false)); + std::vector table_primary_keys = partition_keys; + table_primary_keys.insert(table_primary_keys.end(), primary_keys.begin(), + primary_keys.end()); + ASSERT_OK(catalog->CreateTable(Identifier("foo", "bar"), c_schema.get(), partition_keys, + table_primary_keys, options_, /*ignore_if_exists=*/false)); + } + Result> CreateRealtimeWriter( const std::shared_ptr& realtime_context) const { WriteContextBuilder builder(table_path_, commit_user_); @@ -240,6 +386,12 @@ class RealtimeWriteInteTest : public ::testing::Test { Result> MakeBatch(const std::vector& rows, bool partitioned, int32_t bucket) const { + return MakeBatch(rows, partitioned, bucket, /*row_kinds=*/{}); + } + + Result> MakeBatch( + const std::vector& rows, bool partitioned, int32_t bucket, + const std::vector& row_kinds) const { if (rows.empty()) { return Status::Invalid("cannot create an empty test batch"); } @@ -247,7 +399,7 @@ class RealtimeWriteInteTest : public ::testing::Test { std::string json = "["; for (size_t i = 0; i < rows.size(); ++i) { const auto& [id, payload, pt] = rows[i]; - if (pt != partition) { + if (partitioned && pt != partition) { return Status::Invalid("one test batch must contain only one partition"); } if (i > 0) { @@ -263,6 +415,7 @@ class RealtimeWriteInteTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); RecordBatchBuilder builder(&c_array); + builder.SetRowKinds(row_kinds); if (partitioned) { builder.SetPartition({{"pt", partition}}); } @@ -441,6 +594,27 @@ class RealtimeWriteInteTest : public ::testing::Test { return scan->CreatePlan(); } + Result> CreateQueryReader( + const std::shared_ptr& plan, + const std::shared_ptr& realtime_context) const { + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + return table_read->CreateReader(plan->Splits()); + } + + Result> CreateQueryReader( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + return CreateQueryReader(plan, realtime_context); + } + Result ReadPlan(const std::shared_ptr& plan, const std::shared_ptr& realtime_context, const std::vector& read_fields, @@ -573,6 +747,80 @@ class RealtimeWriteInteTest : public ::testing::Test { return memory_usage; } + Result> ReadPkSequences( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + realtime_context_impl->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected one PK real-time read view"); + } + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_)); + SchemaManager schema_manager(core_options.GetFileSystem(), table_path_); + PAIMON_ASSIGN_OR_RAISE(std::optional> table_schema, + schema_manager.Latest()); + if (!table_schema) { + return Status::Invalid("expected a table schema"); + } + auto read_schema = std::make_unique(); + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(table_schema.value()->Fields()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema( + *RealtimePrimaryKeyLayout::CreateSchema(value_schema->fields()), read_schema.get())); + ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + views[0].store->CreateQueryReaders(views[0].read_view, + /*offset_begin=*/0, query_context)); + std::vector sequences; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr imported, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(imported); + if (!values) { + return Status::Invalid("PK query reader did not return a StructArray"); + } + std::shared_ptr sequence_array = + std::dynamic_pointer_cast( + values->GetFieldByName(SpecialFields::SequenceNumber().Name())); + if (!sequence_array) { + return Status::Invalid("PK query reader did not return sequence numbers"); + } + for (int64_t row = 0; row < sequence_array->length(); ++row) { + sequences.push_back(sequence_array->Value(row)); + } + } + reader->Close(); + } + return sequences; + } + + static std::vector> NewFiles( + const std::vector& progresses) { + std::vector> files; + for (const RealtimeCommitProgress& progress : progresses) { + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + if (!message) { + continue; + } + const std::vector>& new_files = + message->GetNewFilesIncrement().NewFiles(); + files.insert(files.end(), new_files.begin(), new_files.end()); + } + return files; + } + static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { std::vector seen(static_cast(total_rows), false); int64_t max_id = -1; @@ -598,6 +846,8 @@ class RealtimeWriteInteTest : public ::testing::Test { return Status::OK(); } + void RunConcurrencyTest(bool primary_key); + Result ReadCommittedOffsets() const { PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); @@ -620,6 +870,30 @@ class RealtimeWriteInteTest : public ::testing::Test { ASSERT_EQ(expected_rows, actual_rows); } + void ReplayPkWalAndCommit(const std::vector& wal, + const std::vector& row_kinds, + int64_t commit_identifier, + const std::vector& expected_rows) const { + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + ASSERT_EQ(expected_rows, actual_rows); + } + void CheckDropDatePartitionRemovesOffset(bool legacy_partition_name_enabled) { fields_ = {arrow::field("id", arrow::int64()), arrow::field("payload", arrow::utf8()), arrow::field("pt", arrow::date32())}; @@ -627,7 +901,6 @@ class RealtimeWriteInteTest : public ::testing::Test { options_[Options::PARTITION_GENERATE_LEGACY_NAME] = legacy_partition_name_enabled ? "true" : "false"; CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -723,6 +996,848 @@ TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } +TEST_F(RealtimeWriteInteTest, TestPkRead) { + CreatePkTable(); + auto saw_query_predicate = std::make_shared>(false); + auto query_view = std::make_shared>(); + auto factory = + MakeDecoratingFactory(saw_query_predicate, query_view); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}, {1, "new-in-run", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr update_batch, + MakeBatch({Row{1, "new", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(update_batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "new", "p0"}, {2, "two", "p0"}}), memory_rows); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + + std::vector second_rows = {{1, "latest", "p0"}, {2, "gone", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector union_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{1, "latest", "p0"}, {3, "three", "p0"}}), union_rows); + + const std::string expected_payload = "new"; + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"payload", FieldType::STRING, + Literal(FieldType::STRING, expected_payload.data(), expected_payload.size())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_plan, + CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN( + CollectedReadResult filtered_result, + ReadPlan(filtered_plan, realtime_context, {"id", "payload", "pt"}, predicate, + /*enable_predicate_filter=*/true)); + ASSERT_EQ(nullptr, filtered_result.data); + ASSERT_FALSE(saw_query_predicate->load(std::memory_order_acquire)); + filtered_result.reader->Close(); + filtered_result.reader.reset(); + ASSERT_OK(writer->Close()); + writer.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr lifetime_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + table_read->CreateReader(lifetime_plan->Splits())); + ASSERT_FALSE(query_view->expired()); + + std::weak_ptr weak_context = realtime_context; + table_read.reset(); + lifetime_plan.reset(); + realtime_context.reset(); + ASSERT_TRUE(weak_context.expired()); + ASSERT_FALSE(query_view->expired()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch read_batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(read_batch)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_array, + ReadResultCollector::GetArray(std::move(read_batch))); + ASSERT_NE(nullptr, read_array); + read_array.reset(); + reader->Close(); + reader.reset(); + ASSERT_TRUE(query_view->expired()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRealtimeReadOptimizedScanUnsupported) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + const std::vector rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector disk_rows, ReadRows()); + ASSERT_EQ(rows, disk_rows); + + ScanContextBuilder scan_builder(table_path_ + "$ro"); + scan_builder.SetOptions(options_).WithRealtimeContext(realtime_context).WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + Result> scan = TableScan::Create(std::move(scan_context)); + ASSERT_TRUE(scan.status().IsNotImplemented()) << scan.status().ToString(); + ASSERT_NE(std::string::npos, scan.status().ToString().find( + "PK real-time union read does not support read-optimized")); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkDeleteInsertAndPinnedReadsAcrossRefresh) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr delete_batch, + MakeBatch({Row{1, "deleted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(delete_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr insert_batch, + MakeBatch({Row{1, "inserted", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(insert_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr pinned_reader, + table_read->CreateReader(reader_plan->Splits())); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector plan_rows, ReadRows(pinned_plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "inserted", "p0"}}), plan_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader_rows, + ReadResultCollector::CollectResult(pinned_reader.get())); + ASSERT_EQ(1, reader_rows->length()); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkMergeDiskSealedAndActive) { + options_[Options::READ_BATCH_SIZE] = "2"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}, {3, "disk-3", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + }; + int64_t commit_identifier = 0; + for (const std::vector& disk_rows : disk_batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ++commit_identifier; + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "sealed-1", "p0"}, Row{2, "deleted-2", "p0"}, Row{4, "sealed-4", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "active-1", "p0"}, Row{4, "deleted-4", "p0"}, Row{5, "active-5", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"payload", "id"}, /*predicate=*/nullptr, + /*enable_predicate_filter=*/false)); + ASSERT_NE(nullptr, result.data); + ASSERT_GT(result.data->num_chunks(), 1); + for (const std::shared_ptr& chunk : result.data->chunks()) { + ASSERT_LE(chunk->length(), 2); + } + std::shared_ptr result_type = arrow::struct_( + {arrow::field("_VALUE_KIND", arrow::int8()), arrow::field("payload", arrow::utf8()), + arrow::field("id", arrow::int64())}); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, "active-1", 1], + [0, "disk-3", 3], + [0, "active-5", 5], + [0, "disk-10", 10], + [0, "disk-11", 11] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*result.data)) + << result.data->ToString(); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkMergeAllDiskSplitsWithMemory) { + options_[Options::SOURCE_SPLIT_OPEN_FILE_COST] = "1"; + options_[Options::SOURCE_SPLIT_TARGET_SIZE] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + const std::vector> disk_batches = { + {{1, "disk-1", "p0"}, {2, "disk-2", "p0"}}, + {{10, "disk-10", "p0"}, {11, "disk-11", "p0"}}, + {{20, "disk-20", "p0"}, {21, "disk-21", "p0"}}, + }; + for (int64_t commit_identifier = 0; + commit_identifier < static_cast(disk_batches.size()); ++commit_identifier) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(disk_batches[commit_identifier], /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(commit_identifier)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, commit_identifier)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr memory_batch, + MakeBatch({Row{1, "memory-1", "p0"}, Row{10, "deleted-10", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(memory_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr realtime_split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, realtime_split); + ASSERT_EQ(3, realtime_split->DiskSplits().size()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "memory-1", "p0"}, + {2, "disk-2", "p0"}, + {11, "disk-11", "p0"}, + {20, "disk-20", "p0"}, + {21, "disk-21", "p0"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkNestedProjectionAcrossDiskAndMemory) { + const std::shared_ptr projected_b = arrow::field("b", arrow::int64()); + fields_ = { + arrow::field("id", arrow::int64()), + arrow::field("payload", arrow::struct_({arrow::field("a", arrow::int64()), projected_b})), + arrow::field("pt", arrow::utf8()), + }; + schema_ = arrow::schema(fields_); + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + auto make_batch = [&](const std::string& json) -> Result> { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields_), json)); + ArrowArray c_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); + RecordBatchBuilder builder(&c_array); + return builder.SetBucket(0).Finish(); + }; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + make_batch(R"([[1, [101, 1001], "p0"], [2, [102, 1002], "p0"]])")); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr sealed_batch, + make_batch(R"([[1, [201, 2001], "p0"], [3, [203, 2003], "p0"]])")); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr active_batch, + make_batch(R"([[1, [301, 3001], "p0"], [4, [304, null], "p0"]])")); + ASSERT_OK(writer->Write(std::move(active_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + + auto projected_schema = arrow::schema({ + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*projected_schema, c_schema.get()).ok()); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadSchema(std::move(c_schema)) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + table_read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, + ReadResultCollector::CollectResult(reader.get())); + const std::shared_ptr result_type = arrow::struct_({ + arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("payload", arrow::struct_({projected_b})), + arrow::field("id", arrow::int64()), + }); + const std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(result_type, R"([ + [0, [3001], 1], + [0, [1002], 2], + [0, [2003], 3], + [0, [null], 4] + ])") + .ValueOrDie(); + ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) + << actual->ToString(); + reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkKeylessProjection) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({Row{1, "disk", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch({Row{1, "memory", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadPlanWithSchemaAndCheck(plan, realtime_context, + arrow::schema({arrow::field("payload", arrow::utf8())}), R"([ + [0, "memory"] + ])"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestCompositePkKeylessProjection) { + CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({Row{1, "key", "disk"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch({Row{1, "key", "memory"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadPlanWithSchemaAndCheck(plan, realtime_context, + arrow::schema({arrow::field("pt", arrow::utf8())}), R"([ + [0, "memory"] + ])"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { + CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({Row{1, "a", "disk-1a"}, Row{1, "b", "disk-1b"}, + Row{2, "a", "disk-2a"}, Row{3, "c", "disk-3c"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector disk_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, disk_progress.size()); + ASSERT_EQ(OffsetRange(0, 4), disk_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(disk_progress).size()); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(disk_progress, /*commit_identifier=*/0)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr sealed_batch, + MakeBatch({Row{1, "a", "sealed-1a"}, Row{1, "b", "deleted-1b"}, Row{2, "b", "sealed-2b"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(sealed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector sealed_progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, sealed_progress.size()); + ASSERT_EQ(OffsetRange(4, 7), sealed_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(sealed_progress).size()); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr active_batch, + MakeBatch({Row{1, "a", "active-1a"}, Row{1, "c", "active-1c"}, Row{2, "a", "active-2a"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(active_batch))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + ASSERT_FALSE(split->DiskSplits().empty()); + ASSERT_EQ(4, split->CommittedEndOffset()); + ASSERT_EQ(10, split->MemoryEndOffset()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ((std::vector{{1, "a", "active-1a"}, + {1, "c", "active-1c"}, + {2, "a", "active-2a"}, + {2, "b", "sealed-2b"}, + {3, "c", "disk-3c"}}), + actual_rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkWriterHandoff) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector first_rows = { + {0, "value-0", "p0"}, {1, "value-1", "p0"}, {2, "value-2", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch(first_rows, /*partitioned=*/false)); + ASSERT_OK(first_writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, first_progress.size()); + ASSERT_EQ(OffsetRange(0, 3), first_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(first_progress).size()); + ASSERT_EQ(0, NewFiles(first_progress)[0]->min_sequence_number); + ASSERT_EQ(2, NewFiles(first_progress)[0]->max_sequence_number); + ASSERT_OK(first_writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(realtime_context)); + const std::vector second_rows = {{0, "updated-0", "p0"}, {3, "value-3", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch(second_rows, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, second_progress.size()); + ASSERT_EQ(OffsetRange(3, 5), second_progress[0].offset_range); + ASSERT_EQ(1, NewFiles(second_progress).size()); + ASSERT_EQ(3, NewFiles(second_progress)[0]->min_sequence_number); + ASSERT_EQ(4, NewFiles(second_progress)[0]->max_sequence_number); + + first_progress.push_back(std::move(second_progress[0])); + ASSERT_OK(Commit(first_progress, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(realtime_context)); + ASSERT_EQ((std::vector{{0, "updated-0", "p0"}, + {1, "value-1", "p0"}, + {2, "value-2", "p0"}, + {3, "value-3", "p0"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkPartitionBucketRecovery) { + options_[Options::BUCKET] = "2"; + CreatePkTable(/*partition_keys=*/{"pt"}); + const RealtimePartitionBucket p0b0({{"pt", "p0"}}, /*bucket=*/0); + const RealtimePartitionBucket p1b1({{"pt", "p1"}}, /*bucket=*/1); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_first_batch, + MakeBatch({Row{0, "p0-zero", "p0"}, Row{1, "p0-one", "p0"}}, + /*partitioned=*/true, /*bucket=*/0)); + ASSERT_OK(first_writer->Write(std::move(p0_first_batch))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p1_first_batch, + MakeBatch({Row{10, "p1-ten", "p1"}, Row{11, "p1-eleven", "p1"}, Row{12, "p1-twelve", "p1"}}, + /*partitioned=*/true, /*bucket=*/1)); + ASSERT_OK(first_writer->Write(std::move(p1_first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector first_progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(2, first_progress.size()); + std::map first_ranges; + std::map> first_sequences; + for (const RealtimeCommitProgress& progress : first_progress) { + first_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + first_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(0, 2), first_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(0, 3), first_ranges.at(p1b1)); + ASSERT_EQ((std::pair(0, 1)), first_sequences.at(p0b0)); + ASSERT_EQ((std::pair(0, 2)), first_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, + Commit(first_progress, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->RefreshCommittedSnapshot(first_snapshot_id)); + ASSERT_OK(first_writer->Close()); + first_writer.reset(); + first_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p0_second_batch, + MakeBatch({Row{0, "p0-zero-new", "p0"}, Row{2, "p0-two", "p0"}}, + /*partitioned=*/true, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p0_second_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_second_batch, + MakeBatch({Row{10, "p1-ten-deleted", "p1"}, Row{13, "p1-thirteen", "p1"}}, + /*partitioned=*/true, /*bucket=*/1, + {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p1_second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, second_progress.size()); + std::map second_ranges; + std::map> second_sequences; + for (const RealtimeCommitProgress& progress : second_progress) { + second_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + second_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(2, 4), second_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(3, 5), second_ranges.at(p1b1)); + ASSERT_EQ((std::pair(2, 3)), second_sequences.at(p0b0)); + ASSERT_EQ((std::pair(3, 4)), second_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_progress, /*commit_identifier=*/1)); + ASSERT_OK(second_writer->RefreshCommittedSnapshot(second_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(second_context)); + std::sort(actual_rows.begin(), actual_rows.end()); + ASSERT_EQ((std::vector{{0, "p0-zero-new", "p0"}, + {1, "p0-one", "p0"}, + {2, "p0-two", "p0"}, + {11, "p1-eleven", "p1"}, + {12, "p1-twelve", "p1"}, + {13, "p1-thirteen", "p1"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); + + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(2, offsets.size()); + ASSERT_EQ(4, offsets.at(p0b0)); + ASSERT_EQ(5, offsets.at(p1b1)); +} + +TEST_F(RealtimeWriteInteTest, TestPkRecovery) { + CreatePkTable(); + + WriteContextBuilder seed_builder(table_path_, commit_user_); + seed_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_context, seed_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_writer, + FileStoreWrite::Create(std::move(seed_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_batch, + MakeBatch({Row{99, "seed", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(seed_writer->Write(std::move(seed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector> seed_messages, + seed_writer->PrepareCommit(/*wait_compaction=*/false, + /*commit_identifier=*/0)); + CommitContextBuilder seed_commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit_context, + seed_commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr seed_commit, + FileStoreCommit::Create(std::move(seed_commit_context))); + ASSERT_OK(seed_commit->Commit(seed_messages, /*commit_identifier=*/0)); + ASSERT_OK(seed_writer->Close()); + const std::vector mutations = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector mutation_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(first_writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_sequences, ReadPkSequences(first_context)); + ASSERT_EQ((std::vector{1, 2, 3, 4}), memory_sequences); + ASSERT_OK_AND_ASSIGN(std::vector progress, + first_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); + ASSERT_EQ(1, NewFiles(progress).size()); + ASSERT_EQ(2, NewFiles(progress)[0]->min_sequence_number); + ASSERT_EQ(memory_sequences.back(), NewFiles(progress)[0]->max_sequence_number); + ASSERT_EQ(1, NewFiles(progress)[0]->delete_row_count); + ASSERT_OK(Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK(first_writer->Close()); + first_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_replay, ReadRows()); + ASSERT_EQ((std::vector{{1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}), + rows_after_replay); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr restart_batch, + MakeBatch({Row{4, "four", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(second_writer->Write(std::move(restart_batch))); + ASSERT_OK_AND_ASSIGN(std::vector restart_sequences, ReadPkSequences(second_context)); + ASSERT_EQ((std::vector{5}), restart_sequences); + ASSERT_OK_AND_ASSIGN(std::vector restart_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/2)); + ASSERT_EQ(1, restart_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), restart_progress[0].offset_range); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->min_sequence_number); + ASSERT_EQ(5, NewFiles(restart_progress)[0]->max_sequence_number); + ASSERT_OK(second_writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompaction) { + options_[Options::NUM_SORTED_RUNS_COMPACTION_TRIGGER] = "1"; + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + int64_t latest_snapshot_id = -1; + constexpr int64_t kCommitRoundsBeforeCompaction = 4; + std::set committed_file_names; + for (int64_t round = 0; round < kCommitRoundsBeforeCompaction; ++round) { + const bool delete_latest_live_row = round == kCommitRoundsBeforeCompaction - 1; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeBatch( + {Row{delete_latest_live_row ? round - 1 : round, + delete_latest_live_row ? "deleted" : "value-" + std::to_string(round), "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + delete_latest_live_row + ? std::vector{RecordBatch::RowKind::DELETE} + : std::vector{})); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(round)); + ASSERT_EQ(1, progress.size()); + std::shared_ptr message = + std::dynamic_pointer_cast(progress[0].commit_message); + ASSERT_NE(nullptr, message); + ASSERT_TRUE(message->GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, NewFiles(progress).size()); + committed_file_names.insert(NewFiles(progress)[0]->file_name); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(progress, round)); + ASSERT_OK(writer->RefreshCommittedSnapshot(latest_snapshot_id)); + ASSERT_OK_AND_ASSIGN(uint64_t memory_usage, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_EQ(0, memory_usage); + } + WriteContextBuilder compact_builder(table_path_, commit_user_); + compact_builder.SetOptions(options_).WithStreamingMode(true); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_context, compact_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compact_writer, + FileStoreWrite::Create(std::move(compact_context))); + ASSERT_OK(compact_writer->Compact(/*partition=*/{}, /*bucket=*/0, + /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> compact_messages, + compact_writer->PrepareCommit(/*wait_compaction=*/true, /*commit_identifier=*/4)); + ASSERT_EQ(1, compact_messages.size()); + std::shared_ptr compact_message = + std::dynamic_pointer_cast(compact_messages[0]); + ASSERT_NE(nullptr, compact_message); + ASSERT_TRUE(compact_message->GetNewFilesIncrement().IsEmpty()); + ASSERT_EQ(kCommitRoundsBeforeCompaction, + compact_message->GetCompactIncrement().CompactBefore().size()); + std::set compacted_file_names; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactBefore()) { + compacted_file_names.insert(file->file_name); + } + ASSERT_EQ(committed_file_names, compacted_file_names); + ASSERT_FALSE(compact_message->GetCompactIncrement().CompactAfter().empty()); + constexpr int64_t kHistoricalMaxSequenceNumber = kCommitRoundsBeforeCompaction - 1; + int64_t compacted_live_max_sequence_number = -1; + for (const std::shared_ptr& file : + compact_message->GetCompactIncrement().CompactAfter()) { + compacted_live_max_sequence_number = + std::max(compacted_live_max_sequence_number, file->max_sequence_number); + } + ASSERT_LT(compacted_live_max_sequence_number, kHistoricalMaxSequenceNumber); + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK(commit->Commit(compact_messages, /*commit_identifier=*/4)); + ASSERT_OK(compact_writer->Close()); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional compact_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(compact_snapshot); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot->GetCommitKind()); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(4, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector compacted_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}}), compacted_rows); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr fresh_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_writer, + CreateRealtimeWriter(fresh_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr fresh_batch, + MakeBatch({Row{4, "value-4", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(fresh_writer->Write(std::move(fresh_batch))); + ASSERT_OK_AND_ASSIGN(std::vector fresh_sequences, ReadPkSequences(fresh_context)); + ASSERT_EQ((std::vector{compacted_live_max_sequence_number + 1}), fresh_sequences); + ASSERT_LT(fresh_sequences.front(), kHistoricalMaxSequenceNumber); + ASSERT_OK_AND_ASSIGN(std::vector fresh_progress, + fresh_writer->PrepareCommitWithProgress(/*commit_identifier=*/5)); + ASSERT_EQ(1, fresh_progress.size()); + ASSERT_EQ(OffsetRange(4, 5), fresh_progress[0].offset_range); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->min_sequence_number); + ASSERT_EQ(compacted_live_max_sequence_number + 1, + NewFiles(fresh_progress)[0]->max_sequence_number); + ASSERT_OK_AND_ASSIGN(latest_snapshot_id, Commit(fresh_progress, /*commit_identifier=*/5)); + ASSERT_OK(fresh_writer->Close()); + + ASSERT_OK_AND_ASSIGN(offsets, ReadCommittedOffsets()); + ASSERT_EQ(5, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_AND_ASSIGN(std::vector final_rows, ReadRows()); + ASSERT_EQ((std::vector{{0, "value-0", "p0"}, {1, "value-1", "p0"}, {4, "value-4", "p0"}}), + final_rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkMultipleStoredBatchesMergeForQueryAndCommit) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{4, "four", "p0"}, Row{2, "two", "p0"}, Row{1, "one", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr second_batch, + MakeBatch({Row{3, "three", "p0"}, Row{2, "deleted", "p0"}, Row{1, "one-new", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(second_batch))); + + const std::vector expected = {{1, "one-new", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector query_rows, ReadRows(realtime_context)); + ASSERT_EQ(expected, query_rows); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 6), progress[0].offset_range); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector rows, ReadRows()); + ASSERT_EQ(expected, rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { options_[Options::TARGET_FILE_ROW_NUM] = "10"; CreateTable(/*partition_keys=*/{}); @@ -1235,52 +2350,6 @@ TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTick ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestVectorReaderFailurePreservesEarlierSplitTicket) { - CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, - MakeBatch(p0_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p0_batch))); - std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, - MakeBatch(p1_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p1_batch))); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, - CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_EQ(2, plan->Splits().size()); - - std::vector> invalid_splits = plan->Splits(); - std::shared_ptr second_split = - std::dynamic_pointer_cast(invalid_splits[1]); - ASSERT_NE(nullptr, second_split); - std::vector> second_disk_splits = second_split->DiskSplits(); - invalid_splits[1] = std::make_shared( - RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), second_split->Partition(), - second_split->Bucket(), std::move(second_disk_splits), second_split->CommittedEndOffset(), - second_split->MemoryEndOffset(), second_split->OpaqueTicket()); - - ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(options_) - .SetReadFieldNames({"id", "payload", "pt"}) - .WithRealtimeContext(realtime_context) - .WithMemoryPool(pool_); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, - TableRead::Create(std::move(read_context))); - ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), - "unsupported real-time split version"); - - std::vector expected_rows = p0_rows; - expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); - ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); -} - TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -2170,8 +3239,12 @@ TEST_F(RealtimeWriteInteTest, TestReopenRealtimeContextAfterRollback) { ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { - CreateTable(/*partition_keys=*/{}); +void RealtimeWriteInteTest::RunConcurrencyTest(bool primary_key) { + if (primary_key) { + CreatePkTable(); + } else { + CreateTable(/*partition_keys=*/{}); + } ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, RealtimeContext::Create()); ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, @@ -2180,7 +3253,45 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { constexpr int32_t kReadThreadCount = 4; constexpr int64_t kBatchCount = 12; constexpr int64_t kRowsPerBatch = 2; - constexpr int64_t kTotalRows = kBatchCount * kRowsPerBatch; + const int64_t total_rows = kBatchCount * (primary_key ? 3 : kRowsPerBatch); + + std::vector> pk_batches; + std::vector> pk_row_kinds; + std::vector> pk_expected_states(1); + if (primary_key) { + std::map current_rows; + for (int64_t batch_index = 0; batch_index < kBatchCount; ++batch_index) { + const int64_t key = batch_index % 4; + const int64_t deleted_key = (key + 2) % 4; + std::vector rows = {{key, "update-" + std::to_string(batch_index), "p0"}, + {key, "latest-" + std::to_string(batch_index), "p0"}, + {deleted_key, "deleted-" + std::to_string(batch_index), "p0"}}; + pk_batches.push_back(rows); + pk_row_kinds.push_back({batch_index < 4 ? RecordBatch::RowKind::INSERT + : RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE}); + current_rows[key] = rows[1]; + current_rows.erase(deleted_key); + std::vector expected; + for (const auto& [id, row] : current_rows) { + static_cast(id); + expected.push_back(row); + } + pk_expected_states.push_back(std::move(expected)); + } + } + + auto validate_read = [&](const std::vector& rows) { + if (!primary_key) { + return ValidateReadPrefix(rows, total_rows); + } + if (std::find(pk_expected_states.begin(), pk_expected_states.end(), rows) == + pk_expected_states.end()) { + return Status::Invalid("PK real-time read does not match any completed write"); + } + return Status::OK(); + }; std::atomic writer_done{false}; std::atomic prepare_done{false}; @@ -2217,10 +3328,14 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { 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"); + std::vector rows = primary_key + ? pk_batches[static_cast(batch_index)] + : MakeRows(batch_index * kRowsPerBatch, kRowsPerBatch, + /*partition=*/"p0"); Result> batch_result = - MakeBatch(rows, /*partitioned=*/false); + primary_key ? MakeBatch(rows, /*partitioned=*/false, /*bucket=*/0, + pk_row_kinds[static_cast(batch_index)]) + : MakeBatch(rows, /*partitioned=*/false); if (state.RecordErrorIfNotOk(batch_result)) { break; } @@ -2355,7 +3470,7 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { if (state.RecordErrorIfNotOk(result)) { break; } - Status status = ValidateReadPrefix(result.value(), kTotalRows); + Status status = validate_read(result.value()); if (state.RecordErrorIfNotOk(status)) { break; } @@ -2397,16 +3512,28 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { 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)); + if (primary_key) { + ASSERT_EQ(pk_expected_states.back(), final_rows); + } else { + ASSERT_EQ(total_rows, static_cast(final_rows.size())); + ASSERT_OK(ValidateReadPrefix(final_rows, total_rows)); + } ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap committed_offsets, ReadCommittedOffsets()); - ASSERT_EQ(kTotalRows, + ASSERT_EQ(total_rows, committed_offsets.at(RealtimePartitionBucket(/*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, TestConcurrentWritePrepareCommitReadAndRefresh) { + RunConcurrencyTest(/*primary_key=*/false); +} + +TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { + RunConcurrencyTest(/*primary_key=*/true); +} + TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { CreateTable(/*partition_keys=*/{"pt"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context,