diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 4d65743ab..153d524d4 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -26,11 +26,7 @@ namespace paimon { /// Factory for Paimon's default Arrow-backed `RealtimeStore`. class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: - /// Creates an Arrow-backed store for one partition and bucket. - Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) override; + 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..60d1afc39 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -25,8 +25,10 @@ #include #include #include +#include #include +#include "arrow/c/abi.h" #include "paimon/reader/batch_reader.h" #include "paimon/realtime/offset_range.h" #include "paimon/record_batch.h" @@ -41,10 +43,36 @@ namespace paimon { class MemoryPool; class Predicate; -/// A table record batch and its framework-assigned contiguous offset range. +struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig { + StatisticsMode statistics_mode; +}; + +struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig { + /// Primary-key fields after removing partition fields, in comparison order. + std::vector trimmed_primary_keys; +}; + +using RealtimeStoreCreateConfig = + std::variant; + +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 prepared transport schema: + /// [_VALUE_KIND, _SEQUENCE_NUMBER, _REALTIME_OFFSET, table write fields]. + std::unique_ptr<::ArrowSchema> write_schema; + std::map options; + std::shared_ptr memory_pool; + std::map partition; + int32_t bucket = -1; + RealtimeStoreCreateConfig mode_config; +}; + +/// 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 prepared transport schema, are sorted +/// by full primary key then sequence number, and retain the original offset in +/// `_REALTIME_OFFSET`. struct PAIMON_EXPORT RealtimeWriteBatch { /// Input batch whose ownership is transferred to `RealtimeStore::Write`. std::unique_ptr batch; @@ -79,7 +107,10 @@ class PAIMON_EXPORT RealtimeReadView { /// Parameters used by a `RealtimeStore` to create readers for a query. struct PAIMON_EXPORT RealtimeQueryContext { - /// Requested output fields before the mandatory leading `_VALUE_KIND` field is added. + /// Append mode receives the requested output fields before the mandatory leading + /// `_VALUE_KIND` field is added. Primary-key mode receives the complete prepared schema. + /// 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 +147,10 @@ class PAIMON_EXPORT RealtimeStore { /// Creates readers that expose all rows in a sealed segment for Paimon file writing. /// - /// Concatenating the returned readers must produce every sealed row exactly once 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 prepared transport schema; each reader's complete stream is sorted by full + /// primary key then sequence number. virtual Result>> CreateCommitReaders( const std::shared_ptr& segment) = 0; @@ -128,13 +160,15 @@ class PAIMON_EXPORT RealtimeStore { /// also provide a consistent snapshot when a write or seal is in progress. virtual Result> AcquireReadView() = 0; - /// Creates readers over rows in `view` 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 prepared transport schema and may contain multiple mutations per key; each reader's + /// complete stream is sorted by full primary key then sequence number, and the readers + /// collectively expose every raw mutation exactly once. Paimon retains `view` for the lifetime + /// of the resulting framework reader. virtual Result>> CreateQueryReaders( const std::shared_ptr& view, int64_t offset_begin, const RealtimeQueryContext& context) = 0; @@ -157,16 +191,9 @@ class PAIMON_EXPORT RealtimeStoreFactory { public: virtual ~RealtimeStoreFactory() = default; - /// Creates a store configured with the supplied schema, statistics, options, and memory pool. - /// @param write_schema Complete table write schema whose ownership is transferred to the - /// factory. The factory may consume it or retain it in the created store. - /// @param statistics_mode Framework-parsed statistics collection mode. - /// @param options Effective table options available to the store. - /// @param memory_pool Memory pool provided by the write context. - virtual Result> Create( - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool) = 0; + /// Creates a store for the requested table mode and partition-bucket. + /// The factory consumes `request`, including ownership of `request.write_schema`. + virtual Result> Create(RealtimeStoreCreateRequest&& request) = 0; }; } // namespace paimon diff --git a/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 a9810424a..49871672d 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -378,9 +378,12 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/prepared_key_value_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 @@ -780,6 +783,7 @@ 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_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp diff --git a/src/paimon/core/io/merged_key_value_record_reader.cpp b/src/paimon/core/io/merged_key_value_record_reader.cpp index 70f2bcfb9..8c3952874 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader.cpp @@ -117,13 +117,21 @@ Status MergedKeyValueRecordReader::Iterator::LoadNextRawKeyValue() const { } Result> MergedKeyValueRecordReader::NextBatch() { + if (initialization_error_.has_value()) { + return initialization_error_.value(); + } if (visited_) { return std::unique_ptr(); } visited_ = true; auto iterator = std::make_unique(this); - PAIMON_ASSIGN_OR_RAISE(bool has_next, iterator->HasNext()); + Result has_next_result = iterator->HasNext(); + if (!has_next_result.ok()) { + initialization_error_ = has_next_result.status(); + return initialization_error_.value(); + } + bool has_next = std::move(has_next_result).value(); if (!has_next) { return std::unique_ptr(); } diff --git a/src/paimon/core/io/merged_key_value_record_reader.h b/src/paimon/core/io/merged_key_value_record_reader.h index a1b7aa5e4..227a1593a 100644 --- a/src/paimon/core/io/merged_key_value_record_reader.h +++ b/src/paimon/core/io/merged_key_value_record_reader.h @@ -67,6 +67,7 @@ class MergedKeyValueRecordReader : public KeyValueRecordReader { private: bool visited_ = false; + std::optional initialization_error_; std::unique_ptr reader_; std::shared_ptr key_comparator_; std::shared_ptr> merge_function_wrapper_; diff --git a/src/paimon/core/io/merged_key_value_record_reader_test.cpp b/src/paimon/core/io/merged_key_value_record_reader_test.cpp index 1b6b71c69..a0d65205c 100644 --- a/src/paimon/core/io/merged_key_value_record_reader_test.cpp +++ b/src/paimon/core/io/merged_key_value_record_reader_test.cpp @@ -18,8 +18,12 @@ #include "paimon/core/io/merged_key_value_record_reader.h" +#include #include +#include +#include #include +#include #include "arrow/api.h" #include "arrow/array/array_nested.h" @@ -27,17 +31,67 @@ #include "gtest/gtest.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/fields_comparator.h" #include "paimon/core/mergetree/compact/deduplicate_merge_function.h" #include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_fields.h" #include "paimon/memory/memory_pool.h" +#include "paimon/realtime/offset_range.h" #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_key_value_data_file_record_reader.h" #include "paimon/testing/utils/key_value_checker.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/utils/special_field_ids.h" namespace paimon::test { + +namespace { + +std::shared_ptr MakeField(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))); +} + +std::shared_ptr MakePreparedSchema(const arrow::FieldVector& value_fields) { + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), value_fields.begin(), value_fields.end()); + return arrow::schema(prepared_fields); +} + +class TrackingBatchReader : public BatchReader { + public: + TrackingBatchReader(std::unique_ptr&& delegate, int32_t* close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + ++(*close_count_); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + int32_t* close_count_; +}; + +} // namespace + class MergedKeyValueRecordReaderTest : public testing::Test { public: void SetUp() override { @@ -51,6 +105,14 @@ class MergedKeyValueRecordReaderTest : public testing::Test { std::shared_ptr merge_function_wrapper_; }; +TEST_F(MergedKeyValueRecordReaderTest, TestRealtimeOffsetField) { + const DataField& field = RealtimeOffsetField(); + ASSERT_EQ(SpecialFieldIds::REALTIME_OFFSET, field.Id()); + ASSERT_EQ("_REALTIME_OFFSET", field.Name()); + ASSERT_EQ(arrow::Type::INT64, field.Type()->id()); + ASSERT_FALSE(field.Nullable()); +} + TEST_F(MergedKeyValueRecordReaderTest, TestMergeAcrossUnderlyingBatches) { std::vector fields = {DataField(0, arrow::field("k0", arrow::int32())), DataField(1, arrow::field("k1", arrow::int32())), @@ -143,4 +205,372 @@ TEST_F(MergedKeyValueRecordReaderTest, TestSkipMergedNulloptResultInHasNext) { } } +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderOffsetFilter) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 100, 0, 1, 10], + [0, 101, 1, 2, 20], + [0, 102, 4, 3, 30], + [0, 103, 2, 4, 40], + [0, 104, 5, 5, 50], + [0, 105, 3, 6, 60] + ])") + .ValueOrDie()); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 2); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(2, 4), + key_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + std::vector row_kinds = {const_cast(RowKind::Insert()), + const_cast(RowKind::Insert())}; + std::vector levels = {KeyValue::UNKNOWN_LEVEL, KeyValue::UNKNOWN_LEVEL}; + std::vector expected = KeyValueChecker::GenerateKeyValues( + row_kinds, {103, 105}, levels, {{4}, {6}}, {{4, 40}, {6, 60}}, pool_); + KeyValueChecker::CheckResult(expected, results, 1, 2); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestRejectsUnsortedPluginRowsAcrossBatches) { + std::vector value_fields = {DataField(0, arrow::field("id", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + std::shared_ptr prepared_array = + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 2], + [0, 11, 1, 1] + ])") + .ValueOrDie(); + auto batch_reader = + std::make_unique(prepared_array, prepared_type, /*batch_size=*/1); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, + std::nullopt, key_schema, value_schema, pool_)); + Result> result = + ReadResultCollector::CollectKeyValueResult(reader.get()); + ASSERT_NOK_WITH_MSG(result, "not globally sorted by primary key and sequence number"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderCommitSchema) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr extra = MakeField("extra", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, extra}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1, 2]])") + .ValueOrDie()); + + auto query_batch_reader = + std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr query_reader, + AdaptPreparedBatchReader(std::move(query_batch_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector query_results, + (ReadResultCollector::CollectKeyValueResult< + KeyValueRecordReader, KeyValueRecordReader::Iterator>(query_reader.get()))); + ASSERT_EQ(query_results.size(), 1); + ASSERT_EQ(query_results[0].value->GetFieldCount(), 1); + ASSERT_EQ(query_results[0].value->GetInt(0), 1); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); + ASSERT_NOK_WITH_MSG(AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, + std::nullopt, value_schema, value_schema, pool_), + "exact"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestBadCommitBatch) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value = MakeField("value", arrow::int32(), 1); + std::shared_ptr value_schema = arrow::schema({key, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, std::nullopt, + arrow::schema({key}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "field count"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMissingCompositeKey) { + std::shared_ptr key0 = MakeField("key0", arrow::int32(), 0); + std::shared_ptr key1 = MakeField("key1", arrow::int32(), 1); + std::shared_ptr value = MakeField("value", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key0, key1, value}); + std::shared_ptr prepared_schema = MakePreparedSchema({key0, key1, value}); + std::shared_ptr actual_schema = MakePreparedSchema({key0, value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key0, key1}), value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "cannot find field id 1"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestQueryAddRename) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr old_value = MakeField("old_value", arrow::int32(), 1); + std::shared_ptr renamed_value = MakeField("renamed_value", arrow::int32(), 1); + std::shared_ptr added = MakeField("added", arrow::int32(), 2); + std::shared_ptr value_schema = arrow::schema({key, renamed_value, added}); + std::shared_ptr prepared_schema = + MakePreparedSchema({key, renamed_value, added}); + std::shared_ptr actual_schema = MakePreparedSchema({key, old_value}); + std::shared_ptr actual_type = arrow::struct_(actual_schema->fields()); + std::shared_ptr actual = + arrow::ipc::internal::json::ArrayFromJSON(actual_type, R"([[0, 10, 0, 1, 20]])") + .ValueOrDie(); + + auto batch_reader = std::make_unique(actual, actual_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + arrow::schema({key}), value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr iterator, + reader->NextBatch()); + ASSERT_OK_AND_ASSIGN(KeyValue key_value, iterator->Next()); + ASSERT_EQ(20, key_value.value->GetInt(1)); + ASSERT_TRUE(key_value.value->IsNullAt(2)); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestMergedReaderErrorRetry) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + std::shared_ptr prepared_array = + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([[0, 10, 0, 1]])").ValueOrDie(); + auto failing_reader = std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("stable prepared error")); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(failing_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr key_comparator, + FieldsComparator::Create({DataField(0, key)}, true)); + MergedKeyValueRecordReader merged_reader(std::move(reader), key_comparator, + merge_function_wrapper_); + + Result> first = merged_reader.NextBatch(); + Result> retry = merged_reader.NextBatch(); + ASSERT_NOK(first); + ASSERT_NOK(retry); + ASSERT_EQ(first.status().ToString(), retry.status().ToString()); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderSafeDecode) { + std::shared_ptr key = MakeField("key", arrow::int32(), 0); + std::shared_ptr value_schema = arrow::schema({key}); + std::shared_ptr prepared_schema = MakePreparedSchema({key}); + + arrow::FieldVector invalid_fields = prepared_schema->fields(); + invalid_fields[0] = invalid_fields[0]->WithName("wrong_value_kind"); + invalid_fields[3] = MakeField("wrong_key", arrow::int32(), 99); + std::shared_ptr invalid_type = arrow::struct_(invalid_fields); + auto invalid_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(invalid_type, R"([[0, 10, 0, 1]])").ValueOrDie()); + + auto batch_reader = std::make_unique(invalid_array, invalid_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + value_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG( + (ReadResultCollector::CollectKeyValueResult(reader.get())), + "prepared batch field"); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderNestedProjection) { + std::shared_ptr id = MakeField("id", arrow::int32(), 0); + std::shared_ptr item_a = MakeField("a", arrow::int32(), 10); + std::shared_ptr item_b = MakeField("b", arrow::int32(), 11); + std::shared_ptr items = + MakeField("items", arrow::list(arrow::field("item", arrow::struct_({item_a, item_b}))), 2); + std::shared_ptr attr_x = MakeField("x", arrow::int32(), 20); + std::shared_ptr attr_y = MakeField("y", arrow::int32(), 21); + std::shared_ptr attrs = + MakeField("attrs", arrow::map(arrow::utf8(), arrow::struct_({attr_x, attr_y})), 3); + std::shared_ptr key_left = MakeField("left", arrow::int32(), 30); + std::shared_ptr key_right = MakeField("right", arrow::int32(), 31); + std::shared_ptr keyed_values = MakeField( + "keyed_values", arrow::map(arrow::struct_({key_left, key_right}), arrow::int32()), 4); + std::shared_ptr full_value_schema = + arrow::schema({id, items, attrs, keyed_values}); + std::shared_ptr key_schema = arrow::schema({id}); + std::shared_ptr prepared_schema = + MakePreparedSchema(full_value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 9, 9, 0, [[1, 2]], [["prefix", [3, 4]]], [[[5, 6], 7]]], + [0, 10, 0, 1, [[100, 200], [300, 400]], [["k1", [7, 8]], ["k2", [9, 10]]], [[[11, 12], 13], [[21, 22], 23]]], + [0, 11, 11, 2, [[8, 9]], [["suffix", [10, 11]]], [[[12, 13], 14]]] + ])") + .ValueOrDie()); + prepared_array = checked_pointer_cast(prepared_array->Slice(1, 1)); + + std::shared_ptr query_item_b = MakeField("renamed_b", arrow::int32(), 11); + std::shared_ptr query_item_a = MakeField("renamed_a", arrow::int32(), 10); + std::shared_ptr query_items = MakeField( + "items_renamed", + arrow::list(arrow::field("element", arrow::struct_({query_item_b, query_item_a}))), 2); + std::shared_ptr query_attr_y = MakeField("renamed_y", arrow::int32(), 21); + std::shared_ptr query_attr_x = MakeField("renamed_x", arrow::int32(), 20); + std::shared_ptr query_attrs = + MakeField("attrs_renamed", + arrow::map(arrow::utf8(), arrow::struct_({query_attr_y, query_attr_x})), 3); + std::shared_ptr query_key_right = MakeField("renamed_right", arrow::int32(), 31); + std::shared_ptr query_key_left = MakeField("renamed_left", arrow::int32(), 30); + std::shared_ptr query_keyed_values = + MakeField("keyed_values_renamed", + arrow::map(arrow::struct_({query_key_right, query_key_left}), arrow::int32()), 4); + std::shared_ptr query_value_schema = + arrow::schema({id, query_items, query_attrs, query_keyed_values}); + + auto batch_reader = std::make_unique(prepared_array, prepared_type, 1); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(batch_reader), prepared_schema, OffsetRange(0, 1), + key_schema, query_value_schema, pool_)); + ASSERT_OK_AND_ASSIGN( + std::vector results, + (ReadResultCollector::CollectKeyValueResult(reader.get()))); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0].key->GetInt(0), 1); + ASSERT_EQ(results[0].value->GetFieldCount(), 4); + ASSERT_EQ(results[0].value->GetInt(0), 1); + + std::shared_ptr item_array = results[0].value->GetArray(1); + ASSERT_EQ(item_array->Size(), 2); + std::shared_ptr first_item = item_array->GetRow(0, 2); + ASSERT_EQ(first_item->GetInt(0), 200); + ASSERT_EQ(first_item->GetInt(1), 100); + std::shared_ptr second_item = item_array->GetRow(1, 2); + ASSERT_EQ(second_item->GetInt(0), 400); + ASSERT_EQ(second_item->GetInt(1), 300); + + std::shared_ptr attr_map = results[0].value->GetMap(2); + ASSERT_EQ(attr_map->Size(), 2); + std::shared_ptr key_array = attr_map->KeyArray(); + ASSERT_EQ(std::string(key_array->GetStringView(0)), "k1"); + ASSERT_EQ(std::string(key_array->GetStringView(1)), "k2"); + std::shared_ptr value_array = attr_map->ValueArray(); + std::shared_ptr first_attr = value_array->GetRow(0, 2); + ASSERT_EQ(first_attr->GetInt(0), 8); + ASSERT_EQ(first_attr->GetInt(1), 7); + std::shared_ptr second_attr = value_array->GetRow(1, 2); + ASSERT_EQ(second_attr->GetInt(0), 10); + ASSERT_EQ(second_attr->GetInt(1), 9); + + std::shared_ptr keyed_value_map = results[0].value->GetMap(3); + ASSERT_EQ(keyed_value_map->Size(), 2); + std::shared_ptr struct_keys = keyed_value_map->KeyArray(); + std::shared_ptr first_key = struct_keys->GetRow(0, 2); + ASSERT_EQ(first_key->GetInt(0), 12); + ASSERT_EQ(first_key->GetInt(1), 11); + std::shared_ptr second_key = struct_keys->GetRow(1, 2); + ASSERT_EQ(second_key->GetInt(0), 22); + ASSERT_EQ(second_key->GetInt(1), 21); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(0), 13); + ASSERT_EQ(keyed_value_map->ValueArray()->GetInt(1), 23); +} + +TEST_F(MergedKeyValueRecordReaderTest, TestPreparedReaderLifecycle) { + std::vector value_fields = {DataField(0, arrow::field("k0", arrow::int32())), + DataField(1, arrow::field("v0", arrow::int32()))}; + std::shared_ptr value_schema = + DataField::ConvertDataFieldsToArrowSchema(value_fields); + std::shared_ptr key_schema = arrow::schema({value_schema->field(0)}); + std::shared_ptr prepared_schema = MakePreparedSchema(value_schema->fields()); + std::shared_ptr prepared_type = arrow::struct_(prepared_schema->fields()); + auto prepared_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(prepared_type, R"([ + [0, 10, 0, 1, 100] + ])") + .ValueOrDie()); + + int32_t explicit_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &explicit_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + reader->Close(); + } + ASSERT_EQ(explicit_close_count, 1); + + int32_t destructor_close_count = 0; + { + auto tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &destructor_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + } + ASSERT_EQ(destructor_close_count, 1); + + int32_t factory_failure_close_count = 0; + { + std::unique_ptr tracking_reader = std::make_unique( + std::make_unique(prepared_array, prepared_type, 1), + &factory_failure_close_count); + std::shared_ptr invalid_schema = arrow::schema(value_schema->fields()); + ASSERT_NOK(AdaptPreparedBatchReader(std::move(tracking_reader), invalid_schema, + OffsetRange(0, 1), key_schema, value_schema, pool_)); + ASSERT_EQ(nullptr, tracking_reader); + } + ASSERT_EQ(factory_failure_close_count, 1); + + int32_t read_failure_close_count = 0; + { + auto failing_reader = + std::make_unique(prepared_array, prepared_type, 1); + failing_reader->SetNextBatchStatus(Status::IOError("prepared reader failure")); + auto tracking_reader = std::make_unique(std::move(failing_reader), + &read_failure_close_count); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr reader, + AdaptPreparedBatchReader(std::move(tracking_reader), prepared_schema, OffsetRange(0, 1), + key_schema, value_schema, pool_)); + ASSERT_NOK_WITH_MSG(reader->NextBatch(), "prepared reader failure"); + ASSERT_EQ(read_failure_close_count, 1); + } + ASSERT_EQ(read_failure_close_count, 1); +} + } // namespace paimon::test diff --git a/src/paimon/core/mergetree/merge_tree_writer.cpp b/src/paimon/core/mergetree/merge_tree_writer.cpp index 3b6806c73..49961536a 100644 --- a/src/paimon/core/mergetree/merge_tree_writer.cpp +++ b/src/paimon/core/mergetree/merge_tree_writer.cpp @@ -154,6 +154,59 @@ Status MergeTreeWriter::Write(std::unique_ptr&& moved_batch) { return Status::OK(); } +Status MergeTreeWriter::WriteSortedReaders( + std::vector>&& readers) { + auto raw_readers_guard = ScopeGuard([&]() -> void { + for (std::unique_ptr& reader : readers) { + if (reader != nullptr) { + reader->Close(); + } + } + }); + if (readers.empty()) { + return Status::Invalid("sorted readers must not be empty"); + } + for (const std::unique_ptr& reader : readers) { + if (reader == nullptr) { + return Status::Invalid("sorted readers must not contain null reader"); + } + } + + auto sort_merge_reader = std::make_unique( + std::move(readers), key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_); + auto create_consumer = [target_schema = write_schema_, pool = pool_]() + -> Result>> { + return KeyValueMetaProjectionConsumer::Create(target_schema, pool); + }; + auto async_key_value_producer_consumer = + std::make_unique>( + std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), 1, pool_); + raw_readers_guard.Release(); + ScopeGuard async_readers_guard([&]() -> void { async_key_value_producer_consumer->Close(); }); + std::unique_ptr>> rolling_writer; + PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); + ScopeGuard abort_writer_guard([&]() -> void { rolling_writer->Abort(); }); + while (true) { + PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, + async_key_value_producer_consumer->NextBatch()); + if (key_value_batch.batch == nullptr) { + break; + } + PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); + } + PAIMON_RETURN_NOT_OK(rolling_writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, + rolling_writer->GetResult()); + abort_writer_guard.Release(); + + for (const std::shared_ptr& flushed_file : flushed_files) { + new_files_.emplace_back(flushed_file); + PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); + } + metrics_->Merge(rolling_writer->GetMetrics()); + return Status::OK(); +} + Status MergeTreeWriter::Compact(bool full_compaction) { return FlushWriteBuffer(/*wait_for_latest_compaction=*/true, full_compaction); } @@ -256,49 +309,9 @@ Status MergeTreeWriter::FlushWriteBuffer(bool wait_for_latest_compaction, wait_for_latest_compaction = true; } auto cleanup_guard = ScopeGuard([&]() { write_buffer_->Clear(); }); - // 1. flush write buffer to get sorted readers PAIMON_ASSIGN_OR_RAISE(std::vector> readers, write_buffer_->CreateReaders()); - // 2. prepare loser tree sort merge reader - auto sort_merge_reader = std::make_unique( - std::move(readers), key_comparator_, user_defined_seq_comparator_, - merge_function_wrapper_); - // 3. project key value to arrow array - auto create_consumer = [target_schema = write_schema_, pool = pool_]() - -> Result>> { - return KeyValueMetaProjectionConsumer::Create(target_schema, pool); - }; - // consumer batch size is WriteBatchSize - auto async_key_value_producer_consumer = - std::make_unique>( - std::move(sort_merge_reader), create_consumer, options_.GetWriteBatchSize(), - /*projection_thread_num=*/1, pool_); - std::unique_ptr>> - rolling_writer; - PAIMON_ASSIGN_OR_RAISE(rolling_writer, CreateRollingRowWriter()); - ScopeGuard write_guard([&]() -> void { - rolling_writer->Abort(); - async_key_value_producer_consumer->Close(); - }); - while (true) { - PAIMON_ASSIGN_OR_RAISE(KeyValueBatch key_value_batch, - async_key_value_producer_consumer->NextBatch()); - if (key_value_batch.batch == nullptr) { - break; - } - PAIMON_RETURN_NOT_OK(rolling_writer->Write(std::move(key_value_batch))); - } - PAIMON_RETURN_NOT_OK(rolling_writer->Close()); - PAIMON_ASSIGN_OR_RAISE(std::vector> flushed_files, - rolling_writer->GetResult()); - async_key_value_producer_consumer->Close(); - write_guard.Release(); - - for (const auto& flushed_file : flushed_files) { - new_files_.emplace_back(flushed_file); - PAIMON_RETURN_NOT_OK(compact_manager_->AddNewFile(flushed_file)); - } - metrics_->Merge(rolling_writer->GetMetrics()); + PAIMON_RETURN_NOT_OK(WriteSortedReaders(std::move(readers))); } PAIMON_RETURN_NOT_OK(TrySyncLatestCompaction(wait_for_latest_compaction)); PAIMON_RETURN_NOT_OK(compact_manager_->TriggerCompaction(forced_full_compaction)); diff --git a/src/paimon/core/mergetree/merge_tree_writer.h b/src/paimon/core/mergetree/merge_tree_writer.h index febce2afb..01efd975c 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 WriteSortedReaders(std::vector>&& readers); + Status Compact(bool full_compaction) override; Result CompactNotCompleted() override; diff --git a/src/paimon/core/mergetree/merge_tree_writer_test.cpp b/src/paimon/core/mergetree/merge_tree_writer_test.cpp index 2155647a1..9ce5498cb 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 { @@ -211,6 +245,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_; @@ -293,6 +344,29 @@ TEST_P(MergeTreeWriterTest, TestSimple) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [2, 0, "Alice", 10, 0, 13.1], + [0, 0, "Lucy", 20, 1, 14.1], + [1, 0, "Paul", 20, 1, null] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(-1, dir->Str(), sorted_reader_path_factory, 1, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); } TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { @@ -374,6 +448,210 @@ TEST_P(MergeTreeWriterTest, TestWriteMultiBatch) { DataIncrement expected_data_increment({expected_data_file_meta}, /*deleted_files=*/{}, /*changelog_files=*/{}); ASSERT_EQ(expected_data_increment, commit_increment.GetNewFilesIncrement()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [16, 0, "Alice", 10, 0, 113.1], + [14, 0, "Lucy", 20, 1, 114.1], + [13, 0, "Paul", 20, 1, 15.1], + [15, 0, "Skye", 10, 0, 118.1] + ])") + .ValueOrDie()); + auto sorted_reader_path_factory = std::make_shared(); + ASSERT_OK(sorted_reader_path_factory->Init(dir->Str() + "/sorted-readers", "orc", + options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto sorted_reader_writer, + CreateMergeWriter(9, dir->Str(), sorted_reader_path_factory, 0, options)); + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + ASSERT_OK(sorted_reader_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement sorted_reader_commit_increment, + sorted_reader_writer->PrepareCommit(false)); + ASSERT_OK(sorted_reader_writer->Close()); + ASSERT_EQ(1, sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles().size()); + std::string sorted_reader_path = sorted_reader_path_factory->ToPath( + sorted_reader_commit_increment.GetNewFilesIncrement().NewFiles()[0]); + CheckFileContent(sorted_reader_path, expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + std::string uuid = path_factory->uuid_; + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(4, dir->Str(), path_factory, 7, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])") + .ValueOrDie()); + + std::vector> sorted_readers; + sorted_readers.push_back(CreateSingleReader(sorted_reader_array)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + std::string expected_data_file_name = "data-" + uuid + "-0.orc"; + std::string expected_data_file_path = dir->Str() + "/" + expected_data_file_name; + ASSERT_OK_AND_ASSIGN(FileStatus data_file_status, + options.GetFileSystem()->GetFileStatus(expected_data_file_path)); + + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [7, 0, "Alice", 20, 1, 17.1], + [9, 0, "Lucy", 30, 2, 19.1], + [8, 3, "Paul", 10, 3, null] + ])"}, + &expected_array) + .ok()); + CheckFileContent(expected_data_file_path, expected_array); + + ASSERT_TRUE(commit_increment.GetCompactIncrement().IsEmpty()); + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(expected_data_file_name, new_file->file_name); + ASSERT_EQ(data_file_status.GetLen(), new_file->file_size); + ASSERT_EQ(3, new_file->row_count); + ASSERT_EQ(7, new_file->min_sequence_number); + ASSERT_EQ(9, new_file->max_sequence_number); + ASSERT_EQ(7, new_file->schema_id); + ASSERT_EQ(1, new_file->delete_row_count); +} + +TEST_P(MergeTreeWriterTest, TestMergeSortedReaders) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(/*last_sequence_number=*/4, dir->Str(), path_factory, + /*schema_id=*/7, options)); + + auto first_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [5, 0, "Alice", 10, 0, 15.1], + [7, 0, "Carol", 20, 1, 17.1], + [10, 0, "Eve", 30, 2, 20.1] + ])") + .ValueOrDie()); + auto second_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1] + ])") + .ValueOrDie()); + bool first_closed = false; + bool second_closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(first_array), &first_closed)); + sorted_readers.push_back(std::make_unique( + CreateSingleReader(second_array), &second_closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(first_closed); + ASSERT_TRUE(second_closed); + ASSERT_OK_AND_ASSIGN(CommitIncrement commit_increment, merge_writer->PrepareCommit(false)); + ASSERT_OK(merge_writer->Close()); + + ASSERT_EQ(1, commit_increment.GetNewFilesIncrement().NewFiles().size()); + const std::shared_ptr& new_file = + commit_increment.GetNewFilesIncrement().NewFiles()[0]; + ASSERT_EQ(5, new_file->row_count); + ASSERT_EQ(1, new_file->delete_row_count); + std::shared_ptr expected_array; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(write_type_, {R"([ + [5, 0, "Alice", 10, 0, 15.1], + [6, 0, "Bob", 11, 0, 16.1], + [8, 3, "Carol", 21, 1, null], + [9, 0, "David", 22, 2, 19.1], + [10, 0, "Eve", 30, 2, 20.1] + ])"}, + &expected_array) + .ok()); + CheckFileContent(path_factory->ToPath(new_file), expected_array); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderOwnership) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + + bool closed = false; + std::vector> sorted_readers; + sorted_readers.push_back(std::make_unique( + CreateSingleReader(sorted_reader_array), &closed)); + + ASSERT_OK(merge_writer->WriteSortedReaders(std::move(sorted_readers))); + ASSERT_TRUE(closed); + ASSERT_OK(merge_writer->Close()); +} + +TEST_P(MergeTreeWriterTest, TestSortedReaderFailure) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "orc"}})); + + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto path_factory = std::make_shared(); + ASSERT_OK(path_factory->Init(dir->Str(), "orc", options.DataFilePrefix(), nullptr)); + + ASSERT_OK_AND_ASSIGN(auto merge_writer, + CreateMergeWriter(-1, dir->Str(), path_factory, 0, options)); + + std::vector> empty_readers; + Status empty_status = merge_writer->WriteSortedReaders(std::move(empty_readers)); + ASSERT_TRUE(empty_status.IsInvalid()); + + std::vector> null_readers; + null_readers.push_back(nullptr); + Status null_status = merge_writer->WriteSortedReaders(std::move(null_readers)); + ASSERT_TRUE(null_status.IsInvalid()); + + auto sorted_reader_array = std::dynamic_pointer_cast( + arrow::ipc::internal::json::ArrayFromJSON(write_type_, R"([ + [0, 0, "Alice", 10, 0, 13.1] + ])") + .ValueOrDie()); + Status expected_status = Status::IOError("sorted reader failure"); + bool failing_reader_closed = false; + 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->WriteSortedReaders(std::move(failing_readers)); + ASSERT_EQ(expected_status, failing_status); + ASSERT_TRUE(failing_reader_closed); + ASSERT_OK(merge_writer->Close()); } TEST_P(MergeTreeWriterTest, TestSharedShreddingMapDataFileMetaInfo) { diff --git a/src/paimon/core/operation/file_store_write.cpp b/src/paimon/core/operation/file_store_write.cpp index 6807ae35e..84a324762 100644 --- a/src/paimon/core/operation/file_store_write.cpp +++ b/src/paimon/core/operation/file_store_write.cpp @@ -36,6 +36,7 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include "paimon/core/options/merge_engine.h" #include "paimon/core/postpone/postpone_bucket_file_store_write.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" @@ -197,7 +198,26 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRealtimeContext()) { - return Status::Invalid("real-time write currently supports append tables only"); + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(options, *schema)); + if (ignore_previous_files) { + return Status::NotImplemented( + "PK realtime v1 requires restore from the latest snapshot"); + } + if (!ctx->GetWriteSchema().empty()) { + return Status::NotImplemented( + "PK realtime v1 does not support a custom write schema"); + } + PAIMON_ASSIGN_OR_RAISE(std::optional latest_snapshot, + snapshot_manager->LatestSnapshot()); + if (latest_snapshot) { + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetMap realtime_committed_offsets, + RealtimeCommitProperties::ReadOffsets( + latest_snapshot, options.GetFileSystem())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(ctx->GetRealtimeContext())); + PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( + latest_snapshot->Id(), realtime_committed_offsets)); + } } if (options.GetBucket() == BucketModeDefine::POSTPONE_BUCKET) { return PostponeBucketFileStoreWrite::Create( @@ -253,7 +273,8 @@ Result> FileStoreWrite::Create(std::unique_ptrGetRootPath(), schema, arrow_schema, partition_schema, dv_maintainer_factory, io_manager, key_comparator, sequence_fields_comparator, merge_function_wrapper, options, ignore_previous_files, ctx->IsStreamingMode(), ctx->IgnoreNumBucketCheck(), - ctx->EnableMultiThreadSpill(), ctx->GetExecutor(), ctx->GetMemoryPool()); + ctx->EnableMultiThreadSpill(), ctx->GetRealtimeContext(), ctx->GetExecutor(), + ctx->GetMemoryPool()); } } diff --git a/src/paimon/core/operation/key_value_file_store_write.cpp b/src/paimon/core/operation/key_value_file_store_write.cpp index 08c5ea0c3..ee6445057 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,31 @@ #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/common/types/data_field.h" +#include "paimon/core/compact/noop_compact_manager.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #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_fields.h" +#include "paimon/core/realtime/realtime_primary_key_writer.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/utils/file_store_path_factory.h" #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 +70,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 +78,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 +86,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 { @@ -109,19 +125,73 @@ Result> KeyValueFileStoreWrite::CreateWriter( PAIMON_ASSIGN_OR_RAISE( std::shared_ptr levels, Levels::Create(key_comparator_, restore_data_files, options_.GetNumLevels())); - auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr compact_manager, - compact_manager_factory_->CreateCompactManager(partition, bucket, compact_strategy, - compact_executor_, levels, dv_maintainer)); + std::map partition_map; + std::shared_ptr compact_manager; + std::shared_ptr realtime_context_impl; + std::optional realtime_store_state; + if (realtime_context_) { + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + file_store_path_factory_->GeneratePartitionVector(partition)); + partition_map = + std::map(partition_values.begin(), partition_values.end()); + PAIMON_ASSIGN_OR_RAISE(realtime_context_impl, RealtimeContextImpl::Cast(realtime_context_)); + if (schema_->GetFieldByName(RealtimeOffsetField().Name())) { + return Status::Invalid("PK real-time write schema contains reserved transport field " + + RealtimeOffsetField().Name()); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), schema_->fields().begin(), + schema_->fields().end()); + auto c_write_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(std::move(prepared_fields)), c_write_schema.get())); + PAIMON_ASSIGN_OR_RAISE( + RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore(RealtimeStoreCreateRequest{ + std::move(c_write_schema), options_.ToMap(), pool_, partition_map, bucket, + PrimaryKeyRealtimeStoreCreateConfig{trimmed_primary_keys}})); + realtime_store_state = std::move(store_state); + compact_manager = std::make_shared(); + } else { + auto compact_strategy = compact_manager_factory_->CreateCompactStrategy(); + PAIMON_ASSIGN_OR_RAISE(compact_manager, compact_manager_factory_->CreateCompactManager( + partition, bucket, compact_strategy, + compact_executor_, levels, dv_maintainer)); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr writer, MergeTreeWriter::Create( restore_max_seq_number, trimmed_primary_keys, data_file_path_factory, key_comparator_, user_defined_seq_comparator_, merge_function_wrapper_, table_schema_->Id(), schema_, - options_, compact_manager, io_manager_, enable_multi_thread_spill_, pool_)); - return writer; + 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_, trimmed_primary_keys, key_comparator_, + realtime_context_impl, realtime_store_state.value(), restore_max_seq_number, writer, pool_); +} + +Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { + 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..a2344e803 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -19,9 +19,13 @@ #include "paimon/core/operation/key_value_file_store_write.h" #include +#include #include #include +#include +#include #include +#include #include #include @@ -44,6 +48,7 @@ #include "paimon/common/utils/path_util.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/operation/restore_files.h" +#include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/table/sink/commit_message_impl.h" #include "paimon/file_store_commit.h" @@ -52,7 +57,9 @@ #include "paimon/format/file_format_factory.h" #include "paimon/format/reader_builder.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" +#include "paimon/realtime/realtime_context.h" #include "paimon/record_batch.h" #include "paimon/status.h" #include "paimon/testing/utils/test_helper.h" @@ -60,6 +67,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(); +}; + +} class KeyValueFileStoreWriteTest : public ::testing::Test { protected: @@ -127,14 +178,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 +245,58 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { EXPECT_NE(nullptr, metadata); return MapSharedShreddingUtils::DeserializeMetadata(metadata->Copy()).value(); } + + Result>> + ReadPreparedRows(const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + context->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected exactly one real-time store"); + } + RealtimeQueryContext query_context{nullptr, nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + views[0].store->CreateQueryReaders( + views[0].read_view, 0, query_context)); + std::vector> rows; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(array); + if (!values || values->num_fields() != 5) { + return Status::Invalid("unexpected prepared real-time batch"); + } + std::shared_ptr row_kinds = + std::dynamic_pointer_cast(values->field(0)); + std::shared_ptr sequences = + std::dynamic_pointer_cast(values->field(1)); + std::shared_ptr offsets = + std::dynamic_pointer_cast(values->field(2)); + std::shared_ptr ids = + std::dynamic_pointer_cast(values->field(3)); + std::shared_ptr payloads = + std::dynamic_pointer_cast(values->field(4)); + if (!row_kinds || !sequences || !offsets || !ids || !payloads) { + return Status::Invalid("unexpected prepared real-time column type"); + } + for (int64_t row = 0; row < values->length(); ++row) { + rows.emplace_back(row_kinds->Value(row), ids->Value(row), + payloads->GetString(row), sequences->Value(row), + offsets->Value(row)); + } + } + reader->Close(); + } + return rows; + } }; TEST_F(KeyValueFileStoreWriteTest, TestWriteWithInvalidBatch) { @@ -303,6 +407,216 @@ 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 PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ( + (std::vector{{0, 1, "old", 0, 0}, {2, 1, "new", 2, 2}, {3, 2, "two", 1, 1}}), + prepared_rows); + ASSERT_OK_AND_ASSIGN(std::vector progresses, + writer->PrepareCommitWithProgress(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, TestRealtimeOffsetCollision) { + const std::map options = { + {Options::BUCKET, "1"}, {Options::REALTIME_ENABLED, "true"}}; + const std::shared_ptr schema = arrow::schema({ + arrow::field("id", arrow::int64(), false), + arrow::field("_REALTIME_OFFSET", arrow::int64()), + }); + std::unique_ptr dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + CreateTable(dir->Str(), schema, options); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + WriteContextBuilder builder(PathUtil::JoinPath(dir->Str(), "foo.db/bar"), "test"); + builder.SetOptions(options).WithStreamingMode(true).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + FileStoreWrite::Create(std::move(write_context))); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[1, 10]])")), + "PK real-time write schema contains reserved transport field"); + ASSERT_OK(writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, 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 PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector retained_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "one", 0, 0}}), retained_rows); + + std::shared_ptr rejecting_pool = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rejecting_context, + RealtimeContext::Create()); + WriteContextBuilder rejecting_builder(table_path, "rejecting"); + rejecting_builder.SetOptions(options) + .WithStreamingMode(true) + .WithRealtimeContext(rejecting_context) + .WithMemoryPool(rejecting_pool); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_write_context, + rejecting_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr rejecting_writer, + FileStoreWrite::Create(std::move(rejecting_write_context))); + ASSERT_OK(rejecting_writer->Write(MakeBatch(schema, "[]"))); + const int64_t rejecting_allocations_before_write = rejecting_pool->allocation_count; + rejecting_pool->reject_allocations = true; + ASSERT_NOK_WITH_MSG(rejecting_writer->Write(MakeBatch(schema, R"([[2, "two"]])")), + "Out of memory"); + ASSERT_GT(rejecting_pool->allocation_count, rejecting_allocations_before_write); + ASSERT_OK_AND_ASSIGN(std::vector rejected_rows, + ReadPreparedRows(rejecting_context)); + ASSERT_TRUE(rejected_rows.empty()); + ASSERT_OK(rejecting_writer->Close()); +} + +TEST_F(KeyValueFileStoreWriteTest, TestRealtimeLimits) { + const int64_t max = std::numeric_limits::max(); + const std::map options = { + {Options::BUCKET, "1"}, {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 PreparedRow = std::tuple; + ASSERT_OK_AND_ASSIGN(std::vector prepared_rows, + ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + + ASSERT_NOK_WITH_MSG(writer->Write(MakeBatch(schema, R"([[2, "overflow"]])")), + "real-time offset range exceeds INT64_MAX"); + ASSERT_OK_AND_ASSIGN(prepared_rows, ReadPreparedRows(realtime_context)); + ASSERT_EQ((std::vector{{0, 1, "legal", max - 1, max - 1}}), prepared_rows); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context_impl, + RealtimeContextImpl::Cast(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::vector views, + context_impl->AcquireReadViews()); + ASSERT_EQ(1, views.size()); + ASSERT_EQ(std::optional(OffsetRange(max - 1, max)), + views[0].read_view->GetOffsetRange()); + ASSERT_OK(writer->Close()); + ASSERT_GE(snapshot_id, 1); +} + TEST_F(KeyValueFileStoreWriteTest, TestPrepareCommitShouldSucceedWhenDefaultCompactRewriterPathEnabled) { ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index b753ea431..c85e75ee0 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -78,6 +78,126 @@ struct KeyValue; template class MergeFunctionWrapper; +class MergeFileSplitRead::RealtimeReaderBuilder { + public: + static Result> Create( + MergeFileSplitRead* owner, const std::vector>& disk_splits, + std::vector>&& additional_readers) { + 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)); + } + 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::dynamic_pointer_cast(disk_splits.front()); + if (!first_split) { + return Status::Invalid("merge input disk split is not a data split"); + } + const BinaryRow& partition = first_split->Partition(); + const int32_t bucket = first_split->Bucket(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, + owner_->path_factory_->CreateDataFilePathFactory(partition, bucket)); + + std::vector> data_files; + std::vector> deletion_files; + for (const std::shared_ptr& disk_split : disk_splits) { + std::shared_ptr data_split = + std::dynamic_pointer_cast(disk_split); + if (!data_split || !(data_split->Partition() == partition) || + data_split->Bucket() != bucket) { + return Status::Invalid("merge input disk splits do not share a partition-bucket"); + } + if (!data_split->BeforeFiles().empty() || data_split->IsStreaming() || + data_split->Bucket() == BucketModeDefine::POSTPONE_BUCKET) { + return Status::Invalid("additional merge input requires fixed-bucket batch splits"); + } + const std::vector>& split_files = data_split->DataFiles(); + const std::vector>& split_deletion_files = + data_split->DeletionFiles(); + if (!split_deletion_files.empty() && + split_deletion_files.size() != split_files.size()) { + return Status::Invalid( + "merge input disk split deletion files must be empty or match data files"); + } + data_files.insert(data_files.end(), split_files.begin(), split_files.end()); + if (split_deletion_files.empty()) { + deletion_files.insert(deletion_files.end(), split_files.size(), std::nullopt); + } else { + deletion_files.insert(deletion_files.end(), split_deletion_files.begin(), + split_deletion_files.end()); + } + } + + DeletionVector::Factory dv_factory = DeletionVector::CreateFactory( + owner_->options_.GetFileSystem(), + DeletionVector::CreateDeletionFileMap(data_files, deletion_files), owner_->pool_); + std::vector> disk_sections = + IntervalPartition(data_files, owner_->key_comparator_).Partition(); + for (const std::vector& section : disk_sections) { + for (const SortedRun& run : section) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr disk_reader, + owner_->CreateReaderForRun(partition, run, dv_factory, + owner_->predicate_for_keys_, + data_file_path_factory)); + readers->push_back(std::move(disk_reader)); + } + } + return Status::OK(); + } + + Result> CreateMergedReader( + std::vector>&& record_readers) { + if (record_readers.empty()) { + return std::make_unique(std::vector>{}, + owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr sort_merge_reader, + owner_->CreateSortMergeReader(std::move(record_readers))); + return CreateProjectedReader(std::move(sort_merge_reader)); + } + + Result> CreateProjectedReader( + std::unique_ptr&& sort_merge_reader) { + if (!owner_->force_keep_delete_) { + sort_merge_reader = std::make_unique(std::move(sort_merge_reader)); + } + + std::unique_ptr projection_reader; + if (!owner_->context_->EnableMultiThreadRowToBatch()) { + PAIMON_ASSIGN_OR_RAISE( + projection_reader, + KeyValueProjectionReader::Create( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), owner_->pool_)); + } else { + const int32_t thread_number = owner_->context_->GetRowToBatchThreadNumber(); + assert(thread_number > 0); + projection_reader = std::make_unique( + std::move(sort_merge_reader), owner_->raw_read_schema_, owner_->projection_, + owner_->options_.GetReadBatchSize(), thread_number, owner_->pool_); + } + PAIMON_ASSIGN_OR_RAISE(projection_reader, + owner_->ApplyPredicateFilterIfNeeded( + std::move(projection_reader), owner_->context_->GetPredicate())); + return std::make_unique(std::move(projection_reader), + owner_->pool_); + } + + MergeFileSplitRead* owner_; +}; + Result> MergeFileSplitRead::Create( const std::shared_ptr& path_factory, const std::shared_ptr& context, @@ -158,6 +278,12 @@ Result> MergeFileSplitRead::CreateReader( return std::make_unique(std::move(batch_reader), pool_); } +Result> MergeFileSplitRead::CreateRealtimeReader( + const std::vector>& disk_splits, + std::vector>&& additional_readers) { + return RealtimeReaderBuilder::Create(this, disk_splits, std::move(additional_readers)); +} + void MergeFileSplitRead::SetMergeFunctionWrapper( const std::shared_ptr>& merge_function_wrapper) { merge_function_wrapper_ = merge_function_wrapper; diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index d4bfa727c..3cc63a444 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); diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 1d7219c41..babc55a3d 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())); + if (std::holds_alternative(request.mode_config)) { + const AppendRealtimeStoreCreateConfig& append_config = + std::get(request.mode_config); + std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); + return std::make_shared(imported_schema, append_config.statistics_mode, + request.memory_pool, arrow_pool); + } + + const PrimaryKeyRealtimeStoreCreateConfig& config = + std::get(request.mode_config); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create( + imported_schema, config.trimmed_primary_keys, request.memory_pool)); + return std::shared_ptr(std::move(store)); } } // 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..f186a8161 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -232,8 +232,14 @@ TEST_F(ArrowRealtimeStoreTest, TestCommitReaderPreservesSlicedBatch) { TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { ArrowRealtimeStoreFactory factory; std::unique_ptr write_schema = MakeReadSchema(schema_); + RealtimeStoreCreateRequest request{std::move(write_schema), + /*options=*/{}, + pool_, + /*partition=*/{}, + /*bucket=*/0, + AppendRealtimeStoreCreateConfig{StatisticsMode::FULL}}; ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_store, - factory.Create(std::move(write_schema), StatisticsMode::FULL, {}, pool_)); + factory.Create(std::move(request))); std::shared_ptr store = std::dynamic_pointer_cast(realtime_store); ASSERT_NE(nullptr, store); diff --git a/src/paimon/core/realtime/prepared_key_value_reader.cpp b/src/paimon/core/realtime/prepared_key_value_reader.cpp new file mode 100644 index 000000000..864456818 --- /dev/null +++ b/src/paimon/core/realtime/prepared_key_value_reader.cpp @@ -0,0 +1,786 @@ +/* + * 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/prepared_key_value_reader.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_nested.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/buffer.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "arrow/type.h" +#include "arrow/util/bit_util.h" +#include "fmt/format.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/status.h" + +namespace paimon { + +namespace { + +constexpr int32_t kValueKindIndex = 0; +constexpr int32_t kSequenceNumberIndex = 1; +constexpr int32_t kRealtimeOffsetIndex = 2; +constexpr int32_t kPreparedValueStartIndex = 3; + +template +void CloseReaders(const std::vector>& readers) { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } +} + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool); + +class RealtimeOffsetCoverage { + public: + static Result> Create( + const OffsetRange& sealed_offsets, size_t reader_count, + const std::shared_ptr& arrow_pool) { + if (sealed_offsets.begin < 0 || sealed_offsets.end < sealed_offsets.begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr seen_offsets, + arrow::AllocateEmptyBitmap(sealed_offsets.Count(), arrow_pool.get())); + return std::shared_ptr(new RealtimeOffsetCoverage( + sealed_offsets, reader_count, std::move(seen_offsets), arrow_pool)); + } + + Status Add(const arrow::Int64Array& offsets) { + std::lock_guard lock(mutex_); + for (int64_t row = 0; row < offsets.length(); ++row) { + const int64_t offset = offsets.Value(row); + if (offset < sealed_offsets_.begin || offset >= sealed_offsets_.end) { + return Status::Invalid( + "PK real-time store commit reader offset is outside the sealed range"); + } + const int64_t index = offset - sealed_offsets_.begin; + if (arrow::bit_util::GetBit(seen_offsets_->data(), index)) { + return Status::Invalid( + "PK real-time store commit readers contain duplicate REALTIME_OFFSET"); + } + arrow::bit_util::SetBit(seen_offsets_->mutable_data(), index); + ++seen_count_; + } + return Status::OK(); + } + + Status FinishReader() { + std::lock_guard lock(mutex_); + ++finished_reader_count_; + if (finished_reader_count_ == reader_count_ && seen_count_ != sealed_offsets_.Count()) { + return Status::Invalid( + "PK real-time store commit readers did not cover the sealed range"); + } + return Status::OK(); + } + + private: + RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t reader_count, + std::shared_ptr seen_offsets, + const std::shared_ptr& arrow_pool) + : sealed_offsets_(sealed_offsets), + reader_count_(reader_count), + arrow_pool_(arrow_pool), + seen_offsets_(std::move(seen_offsets)) {} + + OffsetRange sealed_offsets_; + size_t reader_count_; + std::shared_ptr arrow_pool_; + std::shared_ptr seen_offsets_; + int64_t seen_count_ = 0; + size_t finished_reader_count_ = 0; + std::mutex mutex_; +}; + +Status CheckPreparedField(const std::shared_ptr& schema, int32_t field_idx, + const DataField& expected_field) { + if (schema->num_fields() <= field_idx) { + return Status::Invalid(fmt::format("prepared schema missing transport field {} at index {}", + expected_field.Name(), field_idx)); + } + const std::shared_ptr& field = schema->field(field_idx); + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, NestedProjectionUtils::GetPaimonFieldId(field)); + if (field->name() != expected_field.Name() || !field->type()->Equals(*expected_field.Type()) || + field->nullable() || field_id != expected_field.Id()) { + return Status::Invalid(fmt::format( + "prepared schema field {} must be non-null {}:{} with field id {}, got {}:{} " + "nullable={} field id {}", + field_idx, expected_field.Name(), expected_field.Type()->ToString(), + expected_field.Id(), field->name(), field->type()->ToString(), field->nullable(), + field_id)); + } + return Status::OK(); +} + +Result FindFieldIndexByPaimonId(const arrow::FieldVector& fields, int32_t field_id) { + std::optional matching_index; + for (int32_t i = 0; i < static_cast(fields.size()); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t candidate_id, + NestedProjectionUtils::GetPaimonFieldId(fields[i])); + if (candidate_id == field_id) { + if (matching_index.has_value()) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + matching_index = i; + } + } + if (matching_index.has_value()) { + return matching_index.value(); + } + return Status::Invalid(fmt::format("cannot find field id {} in prepared schema", field_id)); +} + +Status ValidateProjectionType(const std::shared_ptr& prepared_type, + const std::shared_ptr& query_type) { + if (prepared_type->id() != query_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + switch (query_type->id()) { + case arrow::Type::STRUCT: { + const arrow::FieldVector& prepared_fields = prepared_type->fields(); + for (const std::shared_ptr& query_field : query_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); + } + case arrow::Type::LIST: + return ValidateProjectionType(prepared_type->field(0)->type(), + query_type->field(0)->type()); + case arrow::Type::MAP: { + const std::shared_ptr prepared_map = + checked_pointer_cast(prepared_type); + const std::shared_ptr query_map = + checked_pointer_cast(query_type); + PAIMON_RETURN_NOT_OK( + ValidateProjectionType(prepared_map->key_type(), query_map->key_type())); + return ValidateProjectionType(prepared_map->item_type(), query_map->item_type()); + } + default: + if (!prepared_type->Equals(*query_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + prepared_type->ToString(), query_type->ToString())); + } + return Status::OK(); + } +} + +Status ValidateProjectionSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema) { + arrow::FieldVector prepared_value_fields( + prepared_schema->fields().begin() + kPreparedValueStartIndex, + prepared_schema->fields().end()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx, + FindFieldIndexByPaimonId(prepared_value_fields, query_id)); + PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_value_fields[prepared_idx]->type(), + query_field->type())); + } + return Status::OK(); +} + +Status ValidateExactCommitSchema(const std::shared_ptr& prepared_schema, + const std::shared_ptr& value_schema) { + if (prepared_schema->num_fields() != value_schema->num_fields() + kPreparedValueStartIndex) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + for (int32_t i = 0; i < value_schema->num_fields(); ++i) { + if (!prepared_schema->field(i + kPreparedValueStartIndex) + ->Equals(value_schema->field(i), true)) { + return Status::Invalid("commit requires the exact prepared writer schema"); + } + } + return Status::OK(); +} + +Result> AlignStructArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { + const std::shared_ptr data_type = + checked_pointer_cast(array->type()); + std::unordered_map data_field_id_to_idx; + data_field_id_to_idx.reserve(data_type->num_fields()); + for (int32_t i = 0; i < data_type->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(data_type->field(i))); + if (!data_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared value struct", field_id)); + } + } + + arrow::ArrayVector aligned_arrays; + aligned_arrays.reserve(read_type->num_fields()); + for (const std::shared_ptr& read_field : read_type->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t read_field_id, + NestedProjectionUtils::GetPaimonFieldId(read_field)); + auto data_iter = data_field_id_to_idx.find(read_field_id); + if (data_iter == data_field_id_to_idx.end()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr null_child, + arrow::MakeArrayOfNull(read_field->type(), array->offset() + array->length(), + arrow_pool)); + aligned_arrays.push_back(std::move(null_child)); + continue; + } + std::shared_ptr child = + arrow::MakeArray(array->data()->child_data[data_iter->second]); + PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child, read_field->type(), arrow_pool)); + aligned_arrays.push_back(std::move(child)); + } + + std::shared_ptr aligned_data = array->data()->Copy(); + aligned_data->type = read_type; + aligned_data->child_data.clear(); + aligned_data->child_data.reserve(aligned_arrays.size()); + for (const std::shared_ptr& aligned_array : aligned_arrays) { + aligned_data->child_data.push_back(aligned_array->data()); + } + return arrow::MakeArray(std::move(aligned_data)); +} + +Result> AlignListArrayByPaimonIds( + const std::shared_ptr& array, + const std::shared_ptr& read_type, arrow::MemoryPool* arrow_pool) { + std::shared_ptr values = array->values(); + PAIMON_ASSIGN_OR_RAISE(values, + AlignArrayByPaimonIds(values, read_type->value_type(), arrow_pool)); + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {values->data()}; + return arrow::MakeArray(new_data); +} + +Result> AlignMapArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { + std::shared_ptr keys = array->keys(); + PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys, read_type->key_type(), arrow_pool)); + std::shared_ptr items = array->items(); + PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items, read_type->item_type(), arrow_pool)); + + const std::shared_ptr& entries_data = array->data()->child_data[0]; + std::shared_ptr new_entries = entries_data->Copy(); + new_entries->type = arrow::struct_({read_type->key_field(), read_type->item_field()}); + new_entries->child_data = {keys->data(), items->data()}; + + std::shared_ptr new_data = array->data()->Copy(); + new_data->type = read_type; + new_data->child_data = {std::move(new_entries)}; + return arrow::MakeArray(new_data); +} + +Result> AlignArrayByPaimonIds( + const std::shared_ptr& array, const std::shared_ptr& read_type, + arrow::MemoryPool* arrow_pool) { + if (array->type()->id() != read_type->id()) { + return Status::Invalid(fmt::format("prepared value type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + switch (read_type->id()) { + case arrow::Type::STRUCT: + return AlignStructArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + case arrow::Type::LIST: + return AlignListArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + case arrow::Type::MAP: + return AlignMapArrayByPaimonIds(checked_pointer_cast(array), + checked_pointer_cast(read_type), + arrow_pool); + default: + if (!array->type()->Equals(*read_type)) { + return Status::Invalid( + fmt::format("prepared leaf type {} does not match query type {}", + array->type()->ToString(), read_type->ToString())); + } + return array; + } +} + +Result ProjectFieldsByPaimonIds( + const std::shared_ptr& data_batch, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& query_schema, arrow::MemoryPool* arrow_pool) { + std::unordered_map prepared_field_id_to_idx; + prepared_field_id_to_idx.reserve(prepared_schema->num_fields()); + for (int32_t i = kPreparedValueStartIndex; i < prepared_schema->num_fields(); ++i) { + PAIMON_ASSIGN_OR_RAISE(int32_t field_id, + NestedProjectionUtils::GetPaimonFieldId(prepared_schema->field(i))); + if (!prepared_field_id_to_idx.emplace(field_id, i).second) { + return Status::Invalid( + fmt::format("duplicate field id {} in prepared schema", field_id)); + } + } + + arrow::ArrayVector result; + result.reserve(query_schema->num_fields()); + for (const std::shared_ptr& query_field : query_schema->fields()) { + PAIMON_ASSIGN_OR_RAISE(int32_t query_field_id, + NestedProjectionUtils::GetPaimonFieldId(query_field)); + auto prepared_iter = prepared_field_id_to_idx.find(query_field_id); + if (prepared_iter == prepared_field_id_to_idx.end()) { + return Status::Invalid( + fmt::format("cannot find field id {} in prepared schema", query_field_id)); + } + std::shared_ptr field_array = data_batch->field(prepared_iter->second); + PAIMON_ASSIGN_OR_RAISE(field_array, + AlignArrayByPaimonIds(field_array, query_field->type(), arrow_pool)); + result.push_back(std::move(field_array)); + } + return result; +} + +Result> ApplyOffsetFilter( + const std::shared_ptr& data_batch, + const std::shared_ptr>& offset_array, + const std::optional& visible_offsets, arrow::MemoryPool* arrow_pool) { + if (!visible_offsets.has_value()) { + return data_batch; + } + + arrow::BooleanBuilder filter_builder(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(filter_builder.Reserve(offset_array->length())); + int64_t visible_row_count = 0; + for (int64_t i = 0; i < offset_array->length(); ++i) { + int64_t offset = offset_array->Value(i); + bool visible = offset >= visible_offsets->begin && offset < visible_offsets->end; + filter_builder.UnsafeAppend(visible); + visible_row_count += visible; + } + if (visible_row_count == 0) { + return std::shared_ptr(); + } + if (visible_row_count == data_batch->length()) { + return data_batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr filter, + filter_builder.Finish()); + arrow::compute::ExecContext exec_context(arrow_pool); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum filtered, + arrow::compute::Filter(data_batch, filter, arrow::compute::FilterOptions::Defaults(), + &exec_context)); + return checked_pointer_cast(filtered.make_array()); +} + +class PreparedKeyValueReader final : public KeyValueRecordReader { + public: + PreparedKeyValueReader(std::unique_ptr&& reader, + const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& pool, + const std::shared_ptr& offset_coverage) + : reader_(std::move(reader)), + prepared_schema_(prepared_schema), + visible_offsets_(visible_offsets), + key_schema_(key_schema), + value_schema_(value_schema), + key_comparator_(key_comparator), + pool_(pool), + arrow_pool_(GetArrowPool(pool)), + offset_coverage_(offset_coverage) {} + + ~PreparedKeyValueReader() override { + Close(); + } + + class Iterator final : public KeyValueRecordReader::Iterator { + public: + explicit Iterator(PreparedKeyValueReader* reader) : reader_(reader) {} + + Result HasNext() const override { + return cursor_ < reader_->row_kind_array_->length(); + } + + Result Next() override { + if (cursor_ >= reader_->row_kind_array_->length()) { + return Status::Invalid("No more prepared key values in current iterator"); + } + std::shared_ptr key = + std::make_shared(reader_->key_ctx_, cursor_); + auto value = std::make_unique(reader_->value_ctx_, cursor_); + PAIMON_ASSIGN_OR_RAISE( + const RowKind* row_kind, + RowKind::FromByteValue(reader_->row_kind_array_->Value(cursor_))); + int64_t sequence_number = reader_->sequence_number_array_->Value(cursor_); + ++cursor_; + return KeyValue(row_kind, sequence_number, KeyValue::UNKNOWN_LEVEL, std::move(key), + std::move(value)); + } + + private: + PreparedKeyValueReader* reader_; + int64_t cursor_ = 0; + }; + + Result> NextBatch() override { + if (first_error_.has_value()) { + return first_error_.value(); + } + Result> result = NextBatchImpl(); + if (!result.ok()) { + first_error_ = result.status(); + Close(); + } + return result; + } + + std::shared_ptr GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + if (closed_) { + return; + } + closed_ = true; + ResetBatchState(); + reader_->Close(); + } + + private: + Result> NextBatchImpl() { + while (true) { + ResetBatchState(); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + if (offset_coverage_ && !offset_coverage_finished_) { + offset_coverage_finished_ = true; + PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader()); + } + return std::unique_ptr(); + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!arrow_array || arrow_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("cannot cast prepared batch to StructArray"); + } + std::shared_ptr data_batch = + checked_pointer_cast(arrow_array); + Status transport_status = + ValidatePreparedTransportSchema(arrow::schema(data_batch->type()->fields())); + if (!transport_status.ok()) { + return Status::Invalid( + "prepared batch field does not match prepared transport " + "schema: ", + transport_status.ToString()); + } + if (visible_offsets_.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema( + arrow::schema(data_batch->type()->fields()), key_schema_)); + PAIMON_ASSIGN_OR_RAISE( + arrow_array, + AlignArrayByPaimonIds(data_batch, arrow::struct_(prepared_schema_->fields()), + arrow_pool_.get())); + data_batch = checked_pointer_cast(arrow_array); + } + PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch)); + PAIMON_RETURN_NOT_OK(ValidateOrdering(data_batch)); + + std::shared_ptr> offset_array = + checked_pointer_cast>( + data_batch->field(kRealtimeOffsetIndex)); + if (offset_coverage_) { + PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array)); + } + PAIMON_ASSIGN_OR_RAISE( + data_batch, + ApplyOffsetFilter(data_batch, offset_array, visible_offsets_, arrow_pool_.get())); + if (!data_batch) { + continue; + } + + row_kind_array_ = checked_pointer_cast>( + data_batch->field(kValueKindIndex)); + sequence_number_array_ = checked_pointer_cast>( + data_batch->field(kSequenceNumberIndex)); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + key_schema_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector value_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, + value_schema_, arrow_pool_.get())); + key_ctx_ = std::make_shared(key_fields, pool_); + value_ctx_ = std::make_shared(value_fields, pool_); + ArrowUtils::TraverseArray(data_batch); + return std::make_unique(this); + } + } + + Status ValidatePreparedBatch(const std::shared_ptr& data_batch) const { + if (data_batch->num_fields() != prepared_schema_->num_fields()) { + return Status::Invalid(fmt::format( + "prepared batch field count {} does not match prepared schema field count {}", + data_batch->num_fields(), prepared_schema_->num_fields())); + } + const arrow::FieldVector& batch_fields = data_batch->type()->fields(); + for (int32_t i = 0; i < data_batch->num_fields(); ++i) { + if (!batch_fields[i]->Equals(prepared_schema_->field(i), true)) { + return Status::Invalid(fmt::format( + "prepared batch field {} does not match declared prepared schema", i)); + } + } + if (!data_batch->field(kValueKindIndex) || + data_batch->field(kValueKindIndex)->type_id() != arrow::Type::INT8) { + return Status::Invalid("cannot cast VALUE_KIND column to int8 arrow array"); + } + if (!data_batch->field(kSequenceNumberIndex) || + data_batch->field(kSequenceNumberIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast SEQUENCE_NUMBER column to int64 arrow array"); + } + if (!data_batch->field(kRealtimeOffsetIndex) || + data_batch->field(kRealtimeOffsetIndex)->type_id() != arrow::Type::INT64) { + return Status::Invalid("cannot cast REALTIME_OFFSET column to int64 arrow array"); + } + if (data_batch->field(kValueKindIndex)->null_count() != 0 || + data_batch->field(kSequenceNumberIndex)->null_count() != 0 || + data_batch->field(kRealtimeOffsetIndex)->null_count() != 0) { + return Status::Invalid("prepared transport columns must not contain nulls"); + } + return Status::OK(); + } + + Status ValidateOrdering(const std::shared_ptr& data_batch) { + if (data_batch->length() == 0) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE( + arrow::ArrayVector key_fields, + ProjectFieldsByPaimonIds(data_batch, prepared_schema_, key_schema_, arrow_pool_.get())); + std::shared_ptr key_context = + std::make_shared(key_fields, pool_); + std::shared_ptr sequences = + checked_pointer_cast(data_batch->field(kSequenceNumberIndex)); + for (int64_t row = 0; row < data_batch->length(); ++row) { + ColumnarRowRef current_key(key_context, row); + if (previous_key_context_) { + ColumnarRowRef previous_key(previous_key_context_, previous_key_row_); + const int32_t key_comparison = + key_comparator_->CompareTo(previous_key, current_key); + if (key_comparison > 0 || + (key_comparison == 0 && previous_sequence_ > sequences->Value(row))) { + return Status::Invalid( + "PK real-time plugin reader is not globally sorted by primary key and " + "sequence number"); + } + } + previous_key_context_ = key_context; + previous_key_row_ = row; + previous_sequence_ = sequences->Value(row); + } + return Status::OK(); + } + + void ResetBatchState() { + key_ctx_.reset(); + value_ctx_.reset(); + row_kind_array_.reset(); + sequence_number_array_.reset(); + } + + private: + bool closed_ = false; + std::optional first_error_; + std::unique_ptr reader_; + std::shared_ptr prepared_schema_; + std::optional visible_offsets_; + std::shared_ptr key_schema_; + std::shared_ptr value_schema_; + std::shared_ptr key_comparator_; + std::shared_ptr pool_; + std::shared_ptr arrow_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::shared_ptr previous_key_context_; + int64_t previous_key_row_ = 0; + int64_t previous_sequence_ = 0; +}; + +} // namespace + +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema) { + if (!prepared_schema || prepared_schema->num_fields() < kPreparedValueStartIndex) { + return Status::Invalid("prepared schema must contain realtime transport fields"); + } + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kValueKindIndex, SpecialFields::ValueKind())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kSequenceNumberIndex, SpecialFields::SequenceNumber())); + PAIMON_RETURN_NOT_OK( + CheckPreparedField(prepared_schema, kRealtimeOffsetIndex, RealtimeOffsetField())); + return Status::OK(); +} + +namespace { + +Result> AdaptPreparedBatchReaderImpl( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool, + const std::shared_ptr& offset_coverage) { + std::unique_ptr owned_reader = std::move(reader); + if (!owned_reader) { + return Status::Invalid("prepared batch reader cannot be null"); + } + ScopeGuard close_guard([&owned_reader]() -> void { owned_reader->Close(); }); + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); + if (!key_schema) { + return Status::Invalid("prepared key schema cannot be null"); + } + if (!value_schema) { + return Status::Invalid("prepared value schema cannot be null"); + } + if (!key_comparator) { + return Status::Invalid("prepared key comparator cannot be null"); + } + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, key_schema)); + PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(prepared_schema, value_schema)); + if (!visible_offsets.has_value()) { + PAIMON_RETURN_NOT_OK(ValidateExactCommitSchema(prepared_schema, value_schema)); + } + std::unique_ptr result(new PreparedKeyValueReader( + std::move(owned_reader), prepared_schema, visible_offsets, key_schema, value_schema, + key_comparator, memory_pool, offset_coverage)); + close_guard.Release(); + return result; +} + +} // namespace + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) { + return AdaptPreparedBatchReaderImpl(std::move(reader), prepared_schema, visible_offsets, + key_schema, value_schema, key_comparator, memory_pool, + /*offset_coverage=*/nullptr); +} + +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) { + std::vector> adapted_readers; + ScopeGuard readers_guard([&readers, &adapted_readers]() { + CloseReaders(readers); + CloseReaders(adapted_readers); + }); + if (!memory_pool) { + return Status::Invalid("prepared reader memory pool cannot be null"); + } + for (const std::unique_ptr& reader : readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null commit reader"); + } + } + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr offset_coverage, + RealtimeOffsetCoverage::Create(sealed_offsets, readers.size(), arrow_pool)); + adapted_readers.reserve(readers.size()); + for (std::unique_ptr& reader : readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr adapted_reader, + AdaptPreparedBatchReaderImpl( + std::move(reader), prepared_schema, std::nullopt, key_schema, + value_schema, key_comparator, memory_pool, offset_coverage)); + adapted_readers.push_back(std::move(adapted_reader)); + } + readers_guard.Release(); + return adapted_readers; +} + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool) { + if (!key_schema) { + return Status::Invalid("prepared key schema cannot be null"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector key_fields, + DataField::ConvertArrowSchemaToDataFields(key_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); + return AdaptPreparedBatchReader(std::move(reader), prepared_schema, visible_offsets, key_schema, + value_schema, key_comparator, memory_pool); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/prepared_key_value_reader.h b/src/paimon/core/realtime/prepared_key_value_reader.h new file mode 100644 index 000000000..22a837a76 --- /dev/null +++ b/src/paimon/core/realtime/prepared_key_value_reader.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 +#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 FieldsComparator; +class MemoryPool; + +Status ValidatePreparedTransportSchema(const std::shared_ptr& prepared_schema); + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool); + +Result>> AdaptPreparedCommitBatchReaders( + std::vector>&& readers, + const std::shared_ptr& prepared_schema, const OffsetRange& sealed_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool); + +Result> AdaptPreparedBatchReader( + std::unique_ptr&& reader, const std::shared_ptr& prepared_schema, + const std::optional& visible_offsets, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& memory_pool); + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 000000000..e4c480377 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,509 @@ +/* + * 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 + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "paimon/common/data/columnar/columnar_batch_context.h" +#include "paimon/common/data/columnar/columnar_row_ref.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/core/core_options.h" +#include "paimon/core/index/pk/primary_key_index_definitions.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/macros.h" + +namespace paimon { + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema) { + if (options.GetBucket() <= 0) { + return Status::NotImplemented("PK realtime v1 requires fixed buckets"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime v1 supports only the DEDUPLICATE merge engine"); + } + if (options.DataEvolutionEnabled()) { + return Status::NotImplemented("PK realtime v1 does not support data evolution"); + } + if (!options.GetFieldsSequenceGroups().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence groups"); + } + if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete() || + options.AggregationRemoveRecordOnDelete() || + !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) { + return Status::NotImplemented("PK realtime v1 requires default delete behavior"); + } + if (!options.GetSequenceField().empty()) { + return Status::NotImplemented("PK realtime v1 does not support sequence.field"); + } + if (!options.SequenceFieldSortOrderIsAscending()) { + return Status::NotImplemented( + "PK realtime v1 supports only ascending sequence.field.sort-order"); + } + if (options.NeedLookup() || options.DeletionVectorsEnabled() || + options.GetChangelogProducer() != ChangelogProducer::NONE) { + return Status::NotImplemented("PK realtime v1 does not support lookup or early MOR"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector primary_key_fields, + schema.TrimmedPrimaryKeyFields()); + for (const DataField& field : primary_key_fields) { + if (field.Type()->id() == arrow::Type::FLOAT || field.Type()->id() == arrow::Type::DOUBLE) { + return Status::NotImplemented( + "PK realtime v1 does not support FLOAT or DOUBLE primary keys"); + } + } + if (options.GlobalIndexEnabled()) { + PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions, + PrimaryKeyIndexDefinitions::Create(schema)); + if (!definitions.Definitions().empty()) { + return Status::NotImplemented("PK realtime v1 does not support global indexes"); + } + } + return Status::OK(); +} + +namespace { + +uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t total = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + total += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + total += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + total += GetArrayMemoryUsage(data->dictionary); + } + return total; +} + +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 RawBatchReader final : public BatchReader { + public: + RawBatchReader(std::vector batches, std::vector key_field_indexes, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) + : batches_(std::move(batches)), + positions_(batches_.size(), 0), + key_field_indexes_(std::move(key_field_indexes)), + key_comparator_(key_comparator), + memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + heap_(SourceGreater{this}), + metrics_(std::make_shared()) { + key_contexts_.reserve(batches_.size()); + sequence_arrays_.reserve(batches_.size()); + for (size_t i = 0; i < batches_.size(); ++i) { + const StoredBatch& batch = batches_[i]; + arrow::ArrayVector key_arrays; + key_arrays.reserve(key_field_indexes_.size()); + for (int32_t field_index : key_field_indexes_) { + key_arrays.push_back(batch.data->field(field_index)); + } + key_contexts_.push_back( + std::make_shared(key_arrays, memory_pool_)); + sequence_arrays_.push_back( + checked_pointer_cast(batch.data->field(1))); + if (batch.data->length() > 0) { + heap_.push(i); + } + } + } + + Result NextBatch() override { + if (heap_.empty()) { + return MakeEofBatch(); + } + + struct SelectedRow { + size_t selected_source; + int64_t source_ordinal; + }; + struct SelectedSource { + size_t source; + std::vector rows; + int64_t base = -1; + }; + std::vector selected_rows; + selected_rows.reserve(kOutputBatchSize); + std::vector selected_sources; + std::unordered_map selected_source_indexes; + while (!heap_.empty() && selected_rows.size() < kOutputBatchSize) { + const size_t source = heap_.top(); + heap_.pop(); + auto [source_it, inserted] = + selected_source_indexes.emplace(source, selected_sources.size()); + if (inserted) { + selected_sources.push_back(SelectedSource{source, {}}); + } + SelectedSource& selected_source = selected_sources[source_it->second]; + selected_rows.push_back( + SelectedRow{source_it->second, static_cast(selected_source.rows.size())}); + selected_source.rows.push_back(positions_[source]++); + if (positions_[source] < batches_[source].data->length()) { + heap_.push(source); + } + } + + arrow::compute::ExecContext context(arrow_pool_.get()); + arrow::ArrayVector grouped_batches; + int64_t grouped_row_count = 0; + for (SelectedSource& selected_source : selected_sources) { + arrow::Int64Builder source_index_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + source_index_builder.AppendValues(selected_source.rows)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr source_indices, + source_index_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum source_batch, + arrow::compute::Take(arrow::Datum(batches_[selected_source.source].data), + arrow::Datum(source_indices), + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + selected_source.base = grouped_row_count; + grouped_row_count += static_cast(selected_source.rows.size()); + grouped_batches.push_back(source_batch.make_array()); + } + + std::shared_ptr batch; + if (grouped_batches.size() == 1) { + batch = std::move(grouped_batches[0]); + } else { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr grouped, + arrow::Concatenate(grouped_batches, arrow_pool_.get())); + arrow::Int64Builder order_builder(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(order_builder.Reserve(selected_rows.size())); + for (const SelectedRow& selected : selected_rows) { + order_builder.UnsafeAppend(selected_sources[selected.selected_source].base + + selected.source_ordinal); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr order, + order_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum reordered, + arrow::compute::Take(arrow::Datum(grouped), arrow::Datum(order), + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + batch = reordered.make_array(); + } + auto array = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, array.get(), schema.get())); + return ReadBatch(std::move(array), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + void Close() override { + while (!heap_.empty()) { + heap_.pop(); + } + batches_.clear(); + positions_.clear(); + key_contexts_.clear(); + sequence_arrays_.clear(); + } + + private: + static constexpr size_t kOutputBatchSize = 1024; + + bool Less(size_t left, size_t right) const { + ColumnarRowRef left_key(key_contexts_[left], positions_[left]); + ColumnarRowRef right_key(key_contexts_[right], positions_[right]); + const int32_t key_comparison = key_comparator_->CompareTo(left_key, right_key); + if (key_comparison != 0) { + return key_comparison < 0; + } + const int64_t left_sequence = sequence_arrays_[left]->Value(positions_[left]); + const int64_t right_sequence = sequence_arrays_[right]->Value(positions_[right]); + if (left_sequence != right_sequence) { + return left_sequence < right_sequence; + } + return left < right; + } + + struct SourceGreater { + RawBatchReader* reader; + + bool operator()(size_t left, size_t right) const { + return reader->Less(right, left); + } + }; + + std::vector batches_; + std::vector positions_; + std::vector key_field_indexes_; + std::shared_ptr key_comparator_; + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::vector> key_contexts_; + std::vector> sequence_arrays_; + std::priority_queue, SourceGreater> heap_; + std::shared_ptr metrics_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + Impl(std::shared_ptr prepared_schema, std::vector key_field_indexes, + const std::shared_ptr& key_comparator, + const std::shared_ptr& memory_pool) + : prepared_schema_(std::move(prepared_schema)), + key_field_indexes_(std::move(key_field_indexes)), + key_comparator_(key_comparator), + memory_pool_(memory_pool) {} + + Status Write(RealtimeWriteBatch&& write_batch) { + if (!write_batch.batch || !write_batch.batch->GetData()) { + 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_(prepared_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time prepared batch is not a StructArray"); + } + std::shared_ptr prepared = + checked_pointer_cast(array); + PAIMON_RETURN_NOT_OK_FROM_ARROW(prepared->ValidateFull()); + std::lock_guard lock(mutex_); + if (last_offset_ && write_batch.offset_range.begin != last_offset_.value()) { + return Status::Invalid("PK real-time offset ranges must be contiguous"); + } + building_.push_back( + StoredBatch{prepared, write_batch.offset_range, GetArrayMemoryUsage(prepared->data())}); + building_memory_usage_ += building_.back().memory_usage; + last_offset_ = write_batch.offset_range.end; + 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; + if (!segment->Batches().empty()) { + readers.push_back(std::make_unique( + segment->Batches(), key_field_indexes_, key_comparator_, memory_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::shared_ptr(new ReadView(std::move(segments))); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t, const RealtimeQueryContext&) { + std::shared_ptr typed = std::dynamic_pointer_cast(view); + if (!typed) { + return Status::Invalid("read view was not created by the PK real-time store"); + } + std::vector> readers; + std::vector batches; + for (const std::shared_ptr& segment : typed->Segments()) { + batches.insert(batches.end(), segment->Batches().begin(), segment->Batches().end()); + } + if (!batches.empty()) { + readers.push_back(std::make_unique( + std::move(batches), key_field_indexes_, key_comparator_, memory_pool_)); + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_end) { + std::lock_guard lock(mutex_); + while (!sealed_.empty() && sealed_.front()->GetOffsetRange().end <= committed_end) { + sealed_.erase(sealed_.begin()); + } + return Status::OK(); + } + + uint64_t GetMemoryUsage() const { + std::lock_guard lock(mutex_); + uint64_t 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 prepared_schema_; + std::vector key_field_indexes_; + std::shared_ptr key_comparator_; + std::shared_ptr memory_pool_; + mutable std::mutex mutex_; + std::vector building_; + std::vector> sealed_; + uint64_t building_memory_usage_ = 0; + std::optional last_offset_; +}; + +PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) + : impl_(std::move(impl)) {} +PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; + +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& memory_pool) { + PAIMON_RETURN_NOT_OK(ValidatePreparedTransportSchema(prepared_schema)); + if (trimmed_primary_keys.empty() || !memory_pool) { + return Status::Invalid("PK primary keys are empty or memory pool is null"); + } + std::vector key_field_indexes; + std::vector key_fields; + key_field_indexes.reserve(trimmed_primary_keys.size()); + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + const int32_t field_index = prepared_schema->GetFieldIndex(key); + if (field_index < 3) { + return Status::Invalid("PK field is missing from prepared schema: ", key); + } + key_field_indexes.push_back(field_index); + PAIMON_ASSIGN_OR_RAISE(DataField field, DataField::ConvertArrowFieldToDataField( + prepared_schema->field(field_index))); + key_fields.push_back(std::move(field)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr key_comparator, + FieldsComparator::Create(key_fields, /*is_ascending_order=*/true)); + return std::shared_ptr(new PrimaryKeyRealtimeStore( + std::make_unique(prepared_schema, key_field_indexes, key_comparator, memory_pool))); +} +Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { + return impl_->Write(std::move(batch)); +} +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 offset) { + return impl_->AdvanceCommittedOffset(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..f779b4d7d --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,66 @@ +/* + * 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 CoreOptions; +class MemoryPool; +class TableSchema; + +Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const TableSchema& schema); + +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& memory_pool); + + ~PrimaryKeyRealtimeStore() override; + + Status Write(RealtimeWriteBatch&& batch) override; + Result>> SealForCommit() override; + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override; + Status AdvanceCommittedOffset(int64_t committed_offset) override; + uint64_t GetMemoryUsage() const override; + + private: + class Impl; + explicit PrimaryKeyRealtimeStore(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp new file mode 100644 index 000000000..dc2ce86b1 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,400 @@ +/* + * 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 + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/core/core_options.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr PreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}); +} + +std::shared_ptr NestedPreparedSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField(DataField( + 1, + arrow::field("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}))))}); +} + +std::shared_ptr PkSchema( + const std::shared_ptr& key_type = arrow::int64(), + const std::map& options = {}) { + return TableSchema::Create( + /*schema_id=*/0, + arrow::schema({arrow::field("id", key_type), arrow::field("value", arrow::utf8())}), + /*partition_keys=*/{}, /*primary_keys=*/{"id"}, options) + .value(); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(PreparedSchema()->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +std::unique_ptr MakeBatch(const std::shared_ptr& schema, + const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +void AssertOffsetsZero(const ArrowArray* array) { + ASSERT_NE(nullptr, array); + ASSERT_EQ(0, array->offset); + for (int64_t child = 0; child < array->n_children; ++child) { + AssertOffsetsZero(array->children[child]); + } + if (array->dictionary) { + AssertOffsetsZero(array->dictionary); + } +} + +Result ReadJson(const std::vector>& readers) { + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + batches.push_back(std::move(array)); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches)); + return result->ToString(); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestSupportedOptions) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_OK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestUnsupportedOptions) { + const std::string sequence_group = + std::string(Options::FIELDS_PREFIX) + ".value." + Options::SEQUENCE_GROUP; + const std::vector> unsupported_options = { + {{Options::BUCKET, "0"}}, + {{Options::BUCKET, "1"}, {Options::MERGE_ENGINE, "partial-update"}}, + {{Options::BUCKET, "1"}, {Options::DATA_EVOLUTION_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {sequence_group, "seq"}}, + {{Options::BUCKET, "1"}, {Options::SEQUENCE_FIELD, "seq"}}, + {{Options::BUCKET, "1"}, {Options::FORCE_LOOKUP, "true"}}, + {{Options::BUCKET, "1"}, {Options::DELETION_VECTORS_ENABLED, "true"}}, + {{Options::BUCKET, "1"}, {Options::CHANGELOG_PRODUCER, "input"}}, + }; + for (const std::map& option_map : unsupported_options) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema())); + } +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsFloatingPrimaryKeys) { + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap({{Options::BUCKET, "1"}})); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float32())), + "FLOAT or DOUBLE primary keys"); + ASSERT_NOK_WITH_MSG(ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::float64())), + "FLOAT or DOUBLE primary keys"); +} + +TEST(PrimaryKeyRealtimeStoreOptionsTest, TestRejectsEnabledGlobalIndex) { + const std::map option_map = {{Options::BUCKET, "1"}, + {Options::PK_BTREE_INDEX_COLUMNS, "id"}}; + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(option_map)); + ASSERT_NOK_WITH_MSG( + ValidatePrimaryKeyRealtimeOptions(options, *PkSchema(arrow::int64(), option_map)), + "does not support global indexes"); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, 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_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 3, 3, 3, "three"]])"), + OffsetRange(3, 4)}), + "offset ranges must be contiguous"); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 3, 2, 3, "three"]])"), OffsetRange(2, 3)})); + + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); + ASSERT_GT(store->GetMemoryUsage(), 0); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 4, 4, 4, "four"]])"), OffsetRange(4, 5)}), + "offset ranges must be contiguous"); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestBadTransportPrefix) { + const std::shared_ptr valid = PreparedSchema(); + std::vector invalid_fields; + + arrow::FieldVector wrong_type = valid->fields(); + wrong_type[0] = DataField::ConvertDataFieldToArrowField( + DataField(SpecialFields::ValueKind().Id(), + arrow::field("_VALUE_KIND", arrow::int32(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_type)); + + arrow::FieldVector nullable_sequence = valid->fields(); + nullable_sequence[1] = nullable_sequence[1]->WithNullable(true); + invalid_fields.push_back(std::move(nullable_sequence)); + + arrow::FieldVector wrong_offset_id = valid->fields(); + wrong_offset_id[2] = DataField::ConvertDataFieldToArrowField( + DataField(99, arrow::field("_REALTIME_OFFSET", arrow::int64(), false))) + ->WithNullable(false); + invalid_fields.push_back(std::move(wrong_offset_id)); + + for (const arrow::FieldVector& fields : invalid_fields) { + ASSERT_NOK_WITH_MSG( + PrimaryKeyRealtimeStore::Create(arrow::schema(fields), {"id"}, GetDefaultPool()), + "prepared schema field"); + } +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitBatches) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, 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(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_EQ( + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 2,\n 0\n ]\n-- " + "child 1 type: int64\n [\n 6,\n 7,\n 5\n ]\n-- child 2 type: int64\n [\n " + "1,\n 2,\n 0\n ]\n-- child 3 type: int64\n [\n 1,\n 2,\n 3\n ]\n-- child " + "4 type: string\n [\n \"before\",\n \"after\",\n \"three\"\n ]", + actual); + readers[0]->Close(); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderExportsZeroOffsets) { + std::shared_ptr schema = NestedPreparedSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(schema, {"id"}, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(schema, R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]]])"), + OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + ASSERT_TRUE(arrow::ImportArray(batch.first.get(), batch.second.get()).ok()); + ASSERT_OK_AND_ASSIGN(batch, readers[0]->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestHeapMergeAcrossBatches) { + constexpr int64_t kSourceCount = 2057; + constexpr int64_t kKeyCount = 257; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + for (int64_t source = 0; source < kSourceCount; ++source) { + const int64_t id = (source * 149) % kKeyCount; + const std::string json = + fmt::format(R"([[0, {}, {}, {}, "v{}"]])", source, source, id, source); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(json), OffsetRange(source, source + 1)})); + } + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + + std::vector expected_sources(kSourceCount); + std::iota(expected_sources.begin(), expected_sources.end(), 0); + std::sort(expected_sources.begin(), expected_sources.end(), [=](int64_t left, int64_t right) { + const int64_t left_id = (left * 149) % kKeyCount; + const int64_t right_id = (right * 149) % kKeyCount; + return left_id != right_id ? left_id < right_id : left < right; + }); + + int64_t output_row = 0; + int64_t output_batches = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + ASSERT_LE(batch.first->length, 1024); + ASSERT_GT(batch.first->length, 0); + ++output_batches; + arrow::Result> imported_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(imported_result.ok()) << imported_result.status().ToString(); + std::shared_ptr imported = std::move(imported_result).ValueOrDie(); + std::shared_ptr array = + std::dynamic_pointer_cast(imported); + ASSERT_NE(nullptr, array); + ASSERT_EQ(PreparedSchema()->ToString(), arrow::schema(array->type()->fields())->ToString()); + std::shared_ptr sequences = + std::dynamic_pointer_cast(array->field(1)); + std::shared_ptr ids = + std::dynamic_pointer_cast(array->field(3)); + std::shared_ptr values = + std::dynamic_pointer_cast(array->field(4)); + ASSERT_NE(nullptr, sequences); + ASSERT_NE(nullptr, ids); + ASSERT_NE(nullptr, values); + for (int64_t row = 0; row < array->length(); ++row, ++output_row) { + ASSERT_LT(output_row, kSourceCount); + const int64_t source = expected_sources[output_row]; + ASSERT_EQ(source, sequences->Value(row)); + ASSERT_EQ((source * 149) % kKeyCount, ids->Value(row)); + ASSERT_EQ(fmt::format("v{}", source), values->GetString(row)); + } + } + ASSERT_EQ(kSourceCount, output_row); + ASSERT_EQ(3, output_batches); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCloseUnreadMultiSourceReader) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 10, 0, 1, "a"]])"), OffsetRange(0, 1)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 20, 1, 2, "b"]])"), OffsetRange(1, 2)})); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 30, 2, 3, "c"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + + readers[0]->Close(); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_EQ(std::optional(OffsetRange(4, 5)), view->GetOffsetRange()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderCardinalityIsConstant) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(PreparedSchema(), {"id"}, GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 2, 1, 1, "one"]])"), OffsetRange(1, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + RealtimeQueryContext context{/*read_schema=*/nullptr, /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); +} + +} // namespace +} // namespace paimon::test diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 9d519d791..ea5feecce 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, partition, bucket, + AppendRealtimeStoreCreateConfig{statistics_mode}}; + PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, + realtime_context_impl->GetOrCreateRealtimeStore(std::move(request))); return std::shared_ptr(new RealtimeAppendOnlyWriter( store_state.store, file_writer, input_schema, store_state.initial_offset, memory_pool)); } diff --git a/src/paimon/core/realtime/realtime_context_impl.cpp b/src/paimon/core/realtime/realtime_context_impl.cpp index f6bad5cf1..736ebb02d 100644 --- a/src/paimon/core/realtime/realtime_context_impl.cpp +++ b/src/paimon/core/realtime/realtime_context_impl.cpp @@ -34,14 +34,43 @@ #include #include +#include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/arrow/abi.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/common/utils/uuid.h" #include "paimon/macros.h" #include "paimon/realtime/realtime_store.h" #include "paimon/status.h" namespace paimon { +namespace { + +bool SameMode(const RealtimeStoreCreateConfig& left, const RealtimeStoreCreateConfig& right) { + if (left.index() != right.index()) { + return false; + } + if (const auto* left_pk = std::get_if(&left)) { + const auto& right_pk = std::get(right); + return left_pk->trimmed_primary_keys == right_pk.trimmed_primary_keys; + } + return true; +} + +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 +107,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) { + 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); + const RealtimePartitionBucket key(request.partition, request.bucket); + auto iter = stores_.find(key); int64_t initial_offset = 0; auto offset_iter = committed_offsets_.find(key); if (offset_iter != committed_offsets_.end()) { if (offset_iter->second == std::numeric_limits::max()) { - if (write_schema) { - ArrowSchemaRelease(write_schema.get()); - } 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 (!SameMode(iter->second.mode_config, request.mode_config) || + !iter->second.write_schema->Equals(*requested_schema, /*check_metadata=*/true)) { + return Status::Invalid("real-time store schema or mode mismatch for partition " + + PartitionToString(key.partition) + ", bucket " + + std::to_string(key.bucket) + "; recreate the RealtimeContext"); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - iter->second->AcquireReadView()); + iter->second.store->AcquireReadView()); if (!read_view) { return Status::Invalid("real-time store returned a null read view"); } @@ -117,27 +151,44 @@ 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_ASSIGN_OR_RAISE( - std::shared_ptr store, - factory_->Create(std::move(write_schema), statistics_mode, options, memory_pool)); - stores_.emplace(key, store); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*requested_schema, request.write_schema.get())); + RealtimeStoreCreateConfig mode_config = request.mode_config; + Result> store_result = factory_->Create(std::move(request)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr store, std::move(store_result)); + stores_.emplace(key, + RealtimeStoreRegistryEntry{store, requested_schema, std::move(mode_config)}); if (offset_iter != committed_offsets_.end()) { reclaimed_offsets_.emplace(key, offset_iter->second); } return RealtimeStoreState{std::move(store), initial_offset}; } +int64_t RealtimeContextImpl::AdvanceMaterializedMaxSequenceNumber( + const RealtimePartitionBucket& partition_bucket, int64_t max_sequence_number) { + std::lock_guard lock(mutex_); + auto [iter, inserted] = + materialized_max_sequence_numbers_.emplace(partition_bucket, max_sequence_number); + if (!inserted && max_sequence_number > iter->second) { + iter->second = max_sequence_number; + } + return iter->second; +} + Result> RealtimeContextImpl::AcquireReadViews() { std::lock_guard lock(mutex_); std::vector result; result.reserve(stores_.size()); for (const auto& [partition_bucket, store] : stores_) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_view, - store->AcquireReadView()); + store.store->AcquireReadView()); result.push_back( - RealtimePartitionBucketView{partition_bucket, store, std::move(read_view)}); + RealtimePartitionBucketView{partition_bucket, store.store, std::move(read_view)}); } return result; } @@ -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..f5118c18f 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; @@ -55,6 +59,12 @@ struct RealtimePartitionBucketView { std::shared_ptr read_view; }; +struct RealtimeStoreRegistryEntry { + std::shared_ptr store; + std::shared_ptr write_schema; + RealtimeStoreCreateConfig mode_config; +}; + class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { public: static Result> Create( @@ -65,11 +75,10 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { static Result> Cast( const std::shared_ptr& context); - Result GetOrCreateRealtimeStore( - const std::map& partition, int32_t bucket, - std::unique_ptr<::ArrowSchema> write_schema, StatisticsMode statistics_mode, - const std::map& options, - const std::shared_ptr& memory_pool); + Result GetOrCreateRealtimeStore(RealtimeStoreCreateRequest&& request); + + int64_t AdvanceMaterializedMaxSequenceNumber(const RealtimePartitionBucket& partition_bucket, + int64_t max_sequence_number); Result> AcquireReadViews(); @@ -102,7 +111,8 @@ class PAIMON_EXPORT RealtimeContextImpl final : public RealtimeContext { std::shared_ptr factory_; std::mutex mutex_; std::mutex progress_mutex_; - std::map> stores_; + std::map stores_; + std::map materialized_max_sequence_numbers_; // Full-table progress used as the initial offset when a store is created lazily. RealtimeOffsetMap committed_offsets_; // 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..916b46aad 100644 --- a/src/paimon/core/realtime/realtime_context_test.cpp +++ b/src/paimon/core/realtime/realtime_context_test.cpp @@ -9,12 +9,11 @@ * * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ #include @@ -31,7 +30,6 @@ #include "arrow/c/helpers.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/memory/memory_pool.h" -#include "paimon/realtime/realtime_store.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -49,26 +47,21 @@ class TestingRealtimeStore : public RealtimeStore { Status Write(RealtimeWriteBatch&&) override { return Status::OK(); } - Result>> SealForCommit() override { return std::optional>(); } - Result>> CreateCommitReaders( const std::shared_ptr&) override { return std::vector>(); } - Result> AcquireReadView() override { ++acquire_count; return std::make_shared(); } - Result>> CreateQueryReaders( const std::shared_ptr&, int64_t, const RealtimeQueryContext&) override { return std::vector>(); } - Status AdvanceCommittedOffset(int64_t committed_offset) override { ++advance_count; if (fail_next_advance) { @@ -78,7 +71,6 @@ class TestingRealtimeStore : public RealtimeStore { committed_offsets.push_back(committed_offset); return Status::OK(); } - uint64_t GetMemoryUsage() const override { return 0; } @@ -91,14 +83,11 @@ class TestingRealtimeStore : public RealtimeStore { class TestingRealtimeStoreFactory : public RealtimeStoreFactory { public: - Result> Create(std::unique_ptr write_schema, - StatisticsMode, - const std::map&, - const std::shared_ptr&) override { - if (!write_schema || !write_schema->release) { + Result> Create(RealtimeStoreCreateRequest&& request) override { + if (!request.write_schema || !request.write_schema->release) { return Status::Invalid("testing write schema is null"); } - ArrowSchemaRelease(write_schema.get()); + ArrowSchemaRelease(request.write_schema.get()); auto store = std::make_shared(); stores.push_back(store); return store; @@ -107,12 +96,14 @@ class TestingRealtimeStoreFactory : public RealtimeStoreFactory { 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 +113,39 @@ 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) { + return context->GetOrCreateRealtimeStore( + RealtimeStoreCreateRequest{std::move(write_schema), options, memory_pool, partition, bucket, + AppendRealtimeStoreCreateConfig{StatisticsMode::NONE}}); +} + TEST(RealtimeContextTest, TestReusesStoreAndCapturesRegisteredViews) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState first_state, - 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); + 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 second, + GetOrCreateAppendStore(context, {{"dt", "2026-08-02"}}, 0, + MakeWriteSchema(), {}, GetDefaultPool())); + ASSERT_EQ(first.store, second.store); + ASSERT_EQ(0, second.initial_offset); ASSERT_EQ(1, factory->stores.size()); ASSERT_EQ(1, factory->stores[0]->acquire_count); - ASSERT_OK_AND_ASSIGN( - RealtimeStoreState second_state, - 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 +153,61 @@ 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, 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_EQ(4, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/4)); + ASSERT_EQ(8, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/8)); + ASSERT_EQ(8, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/6)); + ASSERT_EQ(10, context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + /*max_sequence_number=*/10)); +} + TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - 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 +224,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( @@ -211,41 +243,17 @@ TEST(RealtimeContextTest, TestCommittedProgressIsMonotonicAndSelective) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } -TEST(RealtimeContextTest, TestRemovedInactivePartitionDoesNotRequireReopen) { - auto factory = std::make_shared(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - const std::map active_partition = {{"dt", "2026-08-02"}}; - const std::map inactive_partition = {{"dt", "2026-08-03"}}; - const RealtimePartitionBucket active_partition_bucket(active_partition, /*bucket=*/0); - const RealtimePartitionBucket inactive_partition_bucket(inactive_partition, /*bucket=*/0); - - ASSERT_OK(context->AdvanceCommittedProgress( - 5, {{active_partition_bucket, /*offset=*/7}, {inactive_partition_bucket, /*offset=*/9}})); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState active_state, - context->GetOrCreateRealtimeStore(active_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_EQ(7, active_state.initial_offset); - - ASSERT_OK(context->AdvanceCommittedProgress(6, {{active_partition_bucket, /*offset=*/7}})); - ASSERT_OK_AND_ASSIGN(RealtimeStoreState inactive_state, - context->GetOrCreateRealtimeStore(inactive_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_EQ(0, inactive_state.initial_offset); -} - TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); const std::map partition = {{"dt", "2026-08-02"}}; - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 1, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(partition, 2, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); + ASSERT_OK( + GetOrCreateAppendStore(context, partition, 0, MakeWriteSchema(), {}, 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 +267,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)); @@ -271,45 +279,11 @@ TEST(RealtimeContextTest, TestRetriesOnlyIncompleteReclamation) { ASSERT_EQ(std::vector({8}), factory->stores[1]->committed_offsets); } -TEST(RealtimeContextTest, TestRequiresReopenWhenCommittedProgressMovesBackwards) { - auto factory = std::make_shared(); - ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - std::shared_ptr pool = GetDefaultPool(); - const std::map first_partition = {{"dt", "2026-08-02"}}; - const std::map second_partition = {{"dt", "2026-08-03"}}; - const RealtimePartitionBucket first_partition_bucket(first_partition, /*bucket=*/0); - const RealtimePartitionBucket second_partition_bucket(second_partition, /*bucket=*/0); - - ASSERT_OK(context->GetOrCreateRealtimeStore(first_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->GetOrCreateRealtimeStore(second_partition, 0, MakeWriteSchema(), - StatisticsMode::NONE, {}, pool)); - ASSERT_OK(context->AdvanceCommittedProgress( - 5, {{first_partition_bucket, /*offset=*/7}, {second_partition_bucket, /*offset=*/9}})); - ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); - - ASSERT_NOK_WITH_MSG( - context->AdvanceCommittedProgress( - 6, {{first_partition_bucket, /*offset=*/6}, {second_partition_bucket, /*offset=*/10}}), - "recreate RealtimeContext"); - ASSERT_NOK_WITH_MSG( - context->AdvanceCommittedProgress(6, {{first_partition_bucket, /*offset=*/10}}), - "recreate RealtimeContext"); - ASSERT_EQ(std::vector({7}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9}), factory->stores[1]->committed_offsets); - - ASSERT_OK(context->AdvanceCommittedProgress( - 6, {{first_partition_bucket, /*offset=*/10}, {second_partition_bucket, /*offset=*/11}})); - ASSERT_EQ(std::vector({7, 10}), factory->stores[0]->committed_offsets); - ASSERT_EQ(std::vector({9, 11}), factory->stores[1]->committed_offsets); -} - TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); @@ -332,8 +306,8 @@ TEST(RealtimeContextTest, TestPinsResolvesAndReleasesReadViewTicket) { TEST(RealtimeContextTest, TestExpiresAbandonedReadViewTicket) { auto factory = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr context, CreateContext(factory)); - ASSERT_OK(context->GetOrCreateRealtimeStore(/*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), - StatisticsMode::NONE, {}, GetDefaultPool())); + ASSERT_OK(GetOrCreateAppendStore(context, /*partition=*/{}, /*bucket=*/0, MakeWriteSchema(), {}, + GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::vector views, context->AcquireReadViews()); ASSERT_EQ(1, views.size()); diff --git a/src/paimon/core/realtime/realtime_fields.h b/src/paimon/core/realtime/realtime_fields.h new file mode 100644 index 000000000..270941238 --- /dev/null +++ b/src/paimon/core/realtime/realtime_fields.h @@ -0,0 +1,35 @@ +/* + * 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 "arrow/type.h" +#include "paimon/common/types/data_field.h" +#include "paimon/utils/special_field_ids.h" + +namespace paimon { + +inline const DataField& RealtimeOffsetField() { + static const DataField data_field = + DataField(SpecialFieldIds::REALTIME_OFFSET, + arrow::field("_REALTIME_OFFSET", arrow::int64(), /*nullable=*/false)); + return data_field; +} + +} // namespace paimon 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..185156eb7 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,324 @@ +/* + * 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/data_field.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +struct PreparedArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleasePreparedArray(ArrowArray* array) { + auto* data = static_cast(array->private_data); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); + delete data; +} + +Status RetainPreparedArrayPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release || !arrow_pool) { + return Status::Invalid("cannot retain prepared batch memory pool"); + } + array->private_data = + new PreparedArrayPrivateData{array->release, array->private_data, arrow_pool}; + array->release = ReleasePreparedArray; + return Status::OK(); +} + +Result> PrepareBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::vector& trimmed_primary_keys, int64_t first_sequence_number, + int64_t first_offset, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr input, + arrow::ImportArray(batch->GetData(), arrow::struct_(write_schema->fields()))); + if (!input || input->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr values = checked_pointer_cast(input); + const int64_t count = values->length(); + arrow::Int8Builder kinds(arrow_pool); + arrow::Int64Builder sequences(arrow_pool); + arrow::Int64Builder offsets(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count)); + const std::vector& row_kinds = batch->GetRowKind(); + for (int64_t row = 0; row < count; ++row) { + const RecordBatch::RowKind kind = + row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row]; + kinds.UnsafeAppend(static_cast(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr kind_array; + std::shared_ptr sequence_array; + std::shared_ptr offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array)); + arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array), + std::move(offset_array)}; + columns.insert(columns.end(), values->fields().begin(), values->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr prepared, + arrow::StructArray::Make(std::move(columns), prepared_schema->fields())); + + std::vector sort_keys; + sort_keys.reserve(trimmed_primary_keys.size() + 1); + for (const std::string& key : trimmed_primary_keys) { + sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + } + sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), + arrow::compute::SortOrder::Ascending); + arrow::compute::ExecContext context(arrow_pool); + arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum indices, + arrow::compute::SortIndices(arrow::Datum(prepared), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(prepared), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + std::shared_ptr sorted_array = sorted.make_array(); + if (!sorted_array || sorted_array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time sorted batch is not a StructArray"); + } + return checked_pointer_cast(std::move(sorted_array)); +} + +} // namespace + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool) { + if (!store_state.store || !merge_tree_writer || !write_schema || !key_comparator || + !realtime_context || !memory_pool) { + return Status::Invalid("PK real-time writer received a null dependency"); + } + if (trimmed_primary_keys.empty()) { + return Status::Invalid("PK real-time writer requires at least one primary key"); + } + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr field = write_schema->GetFieldByName(key); + if (!field) { + return Status::Invalid("PK field is missing from write schema: ", key); + } + key_fields.push_back(std::move(field)); + } + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), write_schema->fields().begin(), + write_schema->fields().end()); + const RealtimePartitionBucket partition_bucket(partition, bucket); + const int64_t initial_max_sequence_number = + realtime_context->AdvanceMaterializedMaxSequenceNumber(partition_bucket, + restored_max_sequence_number); + return std::shared_ptr(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, + arrow::schema(std::move(prepared_fields)), arrow::schema(std::move(key_fields)), + trimmed_primary_keys, key_comparator, store_state.initial_offset, + initial_max_sequence_number, memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, int64_t next_offset, + int64_t last_sequence_number, const std::shared_ptr& memory_pool) + : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + realtime_store_(realtime_store), + merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), + write_schema_(write_schema), + prepared_schema_(prepared_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t 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 prepared, + PrepareBatch(std::move(batch), write_schema_, prepared_schema_, trimmed_primary_keys_, + first_sequence, next_offset_, arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*prepared, output.get())); + PAIMON_RETURN_NOT_OK(RetainPreparedArrayPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr prepared_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ + count)})); + next_offset_ += count; + last_sequence_number_ += count; + 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> prepared_readers, + AdaptPreparedCommitBatchReaders(std::move(readers), prepared_schema_, sealed_offsets, + key_schema_, write_schema_, key_comparator_, memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(prepared_readers.size()); + for (std::unique_ptr& prepared_reader : prepared_readers) { + auto merge_function = std::make_unique(/*ignore_delete=*/false); + sorted_readers.push_back(std::make_unique( + std::move(prepared_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); + } + return merge_tree_writer_->WriteSortedReaders(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..d65c7e533 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "paimon/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::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool); + + 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& prepared_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + int64_t next_offset, int64_t last_sequence_number, + const std::shared_ptr& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment, + 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 prepared_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; + int64_t next_offset_; + int64_t last_sequence_number_; + std::mutex realtime_store_mutex_; + std::mutex prepare_mutex_; +}; + +} // 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 47603497b..050f09701 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp @@ -46,6 +46,13 @@ TEST(SchemaValidationTest, TestSimple) { ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestRealtimeOffsetIsNotGloballyReserved) { + auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, + TableSchema::Create(0, schema, {}, {}, {})); + ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); +} + TEST(SchemaValidationTest, TestVectorType) { auto vector_field = arrow::field("embedding", arrow::fixed_size_list(arrow::float32(), 3)); auto schema = arrow::schema({arrow::field("id", arrow::int64()), vector_field}); diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index 6885dc374..34c6ef850 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -77,6 +77,13 @@ Result> AppendOnlyTableRead::CreateReader( std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; + ScopeGuard cleanup_guard([&]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); for (const std::shared_ptr& split : splits) { std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); @@ -92,8 +99,6 @@ Result> AppendOnlyTableRead::CreateReader( } } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (!realtime_splits.empty()) { const std::shared_ptr realtime_context = context_->GetRealtimeContext(); if (!realtime_context) { @@ -106,6 +111,9 @@ Result> AppendOnlyTableRead::CreateReader( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + cleanup_guard.Release(); return result; } @@ -124,6 +132,13 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( realtime_context_impl->ResolveReadView(realtime_split->OpaqueTicket())); std::vector> readers; readers.reserve(realtime_split->DiskSplits().size() + 1); + ScopeGuard readers_guard([&readers]() { + for (const std::unique_ptr& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); const RealtimePartitionBucket expected_partition_bucket(realtime_split->Partition(), realtime_split->Bucket()); if (memory.partition_bucket != expected_partition_bucket) { @@ -150,8 +165,17 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( std::vector> memory_readers, memory.store->CreateQueryReaders(memory.read_view, realtime_split->CommittedEndOffset(), query_context)); - + const size_t first_memory_reader = readers.size(); + readers.reserve(readers.size() + memory_readers.size()); for (std::unique_ptr& memory_reader : memory_readers) { + readers.push_back(std::move(memory_reader)); + } + + for (size_t i = first_memory_reader; i < readers.size(); ++i) { + std::unique_ptr& memory_reader = readers[i]; + if (!memory_reader) { + return Status::Invalid("append-only real-time store returned a null query reader"); + } if (context_->EnablePredicateFilter() && context_->GetPredicate()) { PAIMON_ASSIGN_OR_RAISE(memory_reader, PredicateBatchReader::Create( std::move(memory_reader), @@ -159,14 +183,15 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr realtime_reader, RealtimeReader::Create(memory.read_view, std::move(memory_reader))); - readers.push_back(std::move(realtime_reader)); + memory_reader = std::move(realtime_reader); } - std::unique_ptr result = - std::make_unique(std::move(readers), GetMemoryPool()); if (release_ticket) { PAIMON_RETURN_NOT_OK( realtime_context_impl->ReleaseReadView(realtime_split->OpaqueTicket())); } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + readers_guard.Release(); return result; } diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 208807493..3532b59b4 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -20,12 +20,29 @@ #include "paimon/core/table/source/key_value_table_read.h" #include +#include +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/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/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" #include "paimon/core/operation/merge_file_split_read.h" #include "paimon/core/operation/raw_file_split_read.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_fields.h" +#include "paimon/core/realtime/realtime_reader.h" #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/pk_count_reader.h" +#include "paimon/core/table/source/realtime_split.h" #include "paimon/status.h" namespace paimon { @@ -34,6 +51,62 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +struct ColumnarBatchContext; + +namespace { + +Result>> CreateMemoryReaders( + const std::shared_ptr& split, const RealtimePartitionBucketView& memory, + const std::shared_ptr& key_schema, + const std::shared_ptr& value_schema, + const std::shared_ptr& key_comparator, + const std::shared_ptr& context, + const std::shared_ptr& memory_pool) { + std::shared_ptr full_value_schema = + DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields()); + arrow::FieldVector prepared_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(RealtimeOffsetField())->WithNullable(false)}; + prepared_fields.insert(prepared_fields.end(), full_value_schema->fields().begin(), + full_value_schema->fields().end()); + std::shared_ptr prepared_schema = arrow::schema(std::move(prepared_fields)); + auto c_schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema, c_schema.get())); + ScopeGuard schema_guard([schema = c_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{c_schema.get(), nullptr, false}; + PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, + memory.store->CreateQueryReaders(memory.read_view, 0, query_context)); + ScopeGuard batch_readers_guard([&batch_readers]() { + for (const std::unique_ptr& reader : batch_readers) { + if (reader) { + reader->Close(); + } + } + }); + std::vector> result; + result.reserve(batch_readers.size()); + for (std::unique_ptr& reader : batch_readers) { + if (!reader) { + return Status::Invalid("PK real-time store returned a null query reader"); + } + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr prepared_reader, + AdaptPreparedBatchReader( + std::move(reader), prepared_schema, + OffsetRange(split->CommittedEndOffset(), split->MemoryEndOffset()), key_schema, + value_schema, key_comparator, memory_pool)); + auto merge = std::make_unique(false); + result.push_back(std::make_unique( + std::move(prepared_reader), key_comparator, + std::make_shared(std::move(merge)))); + } + batch_readers_guard.Release(); + return result; +} + +} // namespace KeyValueTableRead::KeyValueTableRead(std::vector>&& split_reads, const std::shared_ptr& path_factory, @@ -75,6 +148,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, 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 +204,106 @@ 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, 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())); + } + } + std::unique_ptr result = + std::make_unique(std::move(readers), GetMemoryPool()); + cleanup_guard.Release(); + return result; +} + +Result> KeyValueTableRead::CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket) { + if (realtime_split->Version() != RealtimeSplit::kCurrentVersion) { + return Status::Invalid("unsupported real-time split version"); + } + 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, 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..6824ae59e 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -35,6 +35,7 @@ class Executor; class FileStorePathFactory; class InternalReadContext; class MemoryPool; +class RealtimeSplit; class KeyValueTableRead : public TableRead { public: @@ -45,6 +46,9 @@ class KeyValueTableRead : public TableRead { Result> CreateReader(const std::shared_ptr& split) override; + Result> CreateReader( + const std::vector>& splits) override; + Result> CreateCountReader( const std::vector>& splits) override; @@ -57,6 +61,9 @@ class KeyValueTableRead : public TableRead { const std::shared_ptr& memory_pool, const std::shared_ptr& executor); + Result> CreateRealtimeReader( + const std::shared_ptr& realtime_split, bool release_ticket); + std::vector> split_reads_; std::shared_ptr path_factory_; std::shared_ptr context_; diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index c275208c5..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..f894e1a74 100644 --- a/src/paimon/core/table/source/table_scan.cpp +++ b/src/paimon/core/table/source/table_scan.cpp @@ -41,6 +41,7 @@ #include "paimon/core/operation/data_evolution_file_store_scan.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/operation/key_value_file_store_scan.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/schema_validation.h" @@ -225,15 +226,15 @@ Status ValidateRealtimeScan(const TableSchema& table_schema, const CoreOptions& if (!core_options.RealtimeEnabled()) { return Status::Invalid("real-time scan requires realtime.enabled=true"); } - if (!table_schema.PrimaryKeys().empty()) { - return Status::Invalid("real-time union read currently supports append tables only"); - } if (core_options.GetBucket() <= 0) { return Status::Invalid("real-time union read requires fixed bucket mode"); } if (core_options.DataEvolutionEnabled()) { return Status::Invalid("real-time union read does not support data evolution"); } + if (!table_schema.PrimaryKeys().empty()) { + PAIMON_RETURN_NOT_OK(ValidatePrimaryKeyRealtimeOptions(core_options, table_schema)); + } if (context.IsStreamingMode()) { return Status::Invalid("real-time union read currently supports batch scans only"); } @@ -343,7 +344,7 @@ Result> NewDataTableScan(const std::shared_ptr realtime_context, RealtimeContextImpl::Cast(context->GetRealtimeContext())); return std::make_unique( - std::move(batch_scan), realtime_context, path_factory, + std::move(batch_scan), pk_table, realtime_context, path_factory, snapshot_reader->GetSnapshotManager(), core_options.GetFileSystem(), context->GetScanFilters(), core_options.GetRealtimeReadViewTtlMillis()); } diff --git a/src/paimon/core/utils/primary_key_table_utils_test.cpp b/src/paimon/core/utils/primary_key_table_utils_test.cpp index 12713ca5b..1a7345fdf 100644 --- a/src/paimon/core/utils/primary_key_table_utils_test.cpp +++ b/src/paimon/core/utils/primary_key_table_utils_test.cpp @@ -19,10 +19,11 @@ #include "paimon/core/utils/primary_key_table_utils.h" #include -#include #include #include +#include #include +#include #include "arrow/type.h" #include "gtest/gtest.h" diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 6298137ea..53737fc2b 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -42,7 +42,11 @@ #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/commit_context.h" +#include "paimon/common/factories/io_hook.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/commit/realtime_commit_properties.h" #include "paimon/core/realtime/realtime_context_impl.h" @@ -59,6 +63,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 +76,560 @@ #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 ReadViewCheckingBatchReader final : public BatchReader { + public: + ReadViewCheckingBatchReader(std::unique_ptr delegate, + std::weak_ptr read_view) + : delegate_(std::move(delegate)), read_view_(std::move(read_view)) {} + + Result NextBatch() override { + if (read_view_.expired()) { + return Status::Invalid("real-time read view was released before reader completion"); + } + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::weak_ptr read_view_; +}; + +class QueryTrackingRealtimeStore final : public RealtimeStore { + public: + QueryTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : delegate_(delegate), saw_query_predicate_(saw_query_predicate), query_view_(query_view) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + return delegate_->CreateCommitReaders(segment); + } + + Result> AcquireReadView() override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate_view, + delegate_->AcquireReadView()); + return std::shared_ptr( + std::make_shared(delegate_view)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + if (context.predicate) { + saw_query_predicate_->store(true, std::memory_order_release); + } + *query_view_ = view; + std::shared_ptr tracking_view = + std::dynamic_pointer_cast(view); + if (!tracking_view) { + return Status::Invalid("query tracking store received an unexpected read view"); + } + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + delegate_->CreateQueryReaders(tracking_view->Delegate(), offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), view); + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +class QueryTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit QueryTrackingRealtimeStoreFactory( + const std::shared_ptr>& saw_query_predicate, + const std::shared_ptr>& query_view) + : saw_query_predicate_(saw_query_predicate), query_view_(query_view) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr(std::make_shared( + delegate, saw_query_predicate_, query_view_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr> saw_query_predicate_; + std::shared_ptr> query_view_; +}; + +class CloseTrackingBatchReader final : public BatchReader { + public: + CloseTrackingBatchReader(std::unique_ptr delegate, + const std::shared_ptr>& close_count) + : delegate_(std::move(delegate)), close_count_(close_count) {} + + Result NextBatch() override { + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + close_count_->fetch_add(1, std::memory_order_release); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr> close_count_; +}; + +struct CloseTrackingReaderState { + std::shared_ptr> query_close_count = + std::make_shared>(0); + std::shared_ptr> commit_close_count = + std::make_shared>(0); + int32_t query_null_index = -1; + int32_t commit_null_index = -1; +}; + +class CloseTrackingRealtimeStore final : public RealtimeStore { + public: + CloseTrackingRealtimeStore(const std::shared_ptr& delegate, + const std::shared_ptr& state) + : delegate_(delegate), state_(state) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), + state_->commit_close_count); + } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->commit_null_index, &readers)); + return readers; + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateQueryReaders(view, offset_begin, context)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader), + state_->query_close_count); + } + PAIMON_RETURN_NOT_OK(InsertNullReader(state_->query_null_index, &readers)); + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + static Status InsertNullReader(int32_t index, + std::vector>* readers) { + if (index < 0) { + return Status::OK(); + } + if (index > static_cast(readers->size())) { + return Status::Invalid("null reader index exceeds reader count"); + } + readers->insert(readers->begin() + index, nullptr); + return Status::OK(); + } + + std::shared_ptr delegate_; + std::shared_ptr state_; +}; + +class CloseTrackingRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit CloseTrackingRealtimeStoreFactory( + const std::shared_ptr& state) + : state_(state) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate, state_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + std::shared_ptr state_; +}; + +class SplitBatchReader final : public BatchReader { + public: + explicit SplitBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + while (!current_batch_ || next_row_ == current_batch_->length()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("split batch reader received a non-struct batch"); + } + current_batch_ = std::dynamic_pointer_cast(array); + next_row_ = 0; + } + std::shared_ptr slice = current_batch_->Slice(next_row_, /*length=*/1); + ++next_row_; + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*slice, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + current_batch_.reset(); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::shared_ptr current_batch_; + int64_t next_row_ = 0; +}; + +class SplitCommitReaderRealtimeStore final : public RealtimeStore { + public: + explicit SplitCommitReaderRealtimeStore(const std::shared_ptr& delegate) + : delegate_(delegate) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + reader = std::make_unique(std::move(reader)); + } + return readers; + } + + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; +}; + +class SplitCommitReaderRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate)); + } + + private: + ArrowRealtimeStoreFactory delegate_; +}; + +class DropLastBatchReader final : public BatchReader { + public: + explicit DropLastBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + if (!buffered_.has_value()) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch first, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(first)) { + return MakeEofBatch(); + } + buffered_ = std::move(first); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch next, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(next)) { + buffered_.reset(); + return MakeEofBatch(); + } + ReadBatch result = std::move(buffered_.value()); + buffered_ = std::move(next); + return result; + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + buffered_.reset(); + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; + std::optional buffered_; +}; + +class SwapFirstTwoBatchReader final : public BatchReader { + public: + explicit SwapFirstTwoBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + if (!initialized_) { + initialized_ = true; + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (array->length() < 2) { + return Status::Invalid("cannot make a one-row reader unsorted"); + } + arrow::ArrayVector pieces = {array->Slice(1, 1), array->Slice(0, 1)}; + if (array->length() > 2) { + pieces.push_back(array->Slice(2)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr swapped, + arrow::Concatenate(pieces)); + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*swapped, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + return delegate_->NextBatch(); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + bool initialized_ = false; + std::unique_ptr delegate_; +}; + +class SubstituteOffsetBatchReader final : public BatchReader { + public: + explicit SubstituteOffsetBatchReader(std::unique_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, delegate_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT || array->length() == 0) { + return Status::Invalid("offset substitution requires a non-empty struct batch"); + } + std::shared_ptr struct_array = + std::dynamic_pointer_cast(array); + std::shared_ptr offsets = + std::dynamic_pointer_cast(struct_array->field(2)); + if (!offsets) { + return Status::Invalid("offset substitution requires an int64 REALTIME_OFFSET"); + } + arrow::Int64Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(offsets->length())); + for (int64_t row = 0; row < offsets->length(); ++row) { + builder.UnsafeAppend(0); + } + std::shared_ptr substituted_offsets; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&substituted_offsets)); + std::shared_ptr substituted_data = struct_array->data()->Copy(); + substituted_data->child_data[2] = substituted_offsets->data(); + std::shared_ptr substituted = arrow::MakeArray(std::move(substituted_data)); + auto output = std::make_unique(); + auto schema = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*substituted, output.get(), schema.get())); + return ReadBatch(std::move(output), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return delegate_->GetReaderMetrics(); + } + + void Close() override { + delegate_->Close(); + } + + private: + std::unique_ptr delegate_; +}; + +enum class CommitReaderMalformation { DROP_LAST, UNSORTED, SUBSTITUTE_OFFSET }; + +class MalformedCoverageRealtimeStore final : public RealtimeStore { + public: + MalformedCoverageRealtimeStore(const std::shared_ptr& delegate, + CommitReaderMalformation malformation) + : delegate_(delegate), malformation_(malformation) {} + + Status Write(RealtimeWriteBatch&& batch) override { + return delegate_->Write(std::move(batch)); + } + Result>> SealForCommit() override { + return delegate_->SealForCommit(); + } + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + delegate_->CreateCommitReaders(segment)); + for (std::unique_ptr& reader : readers) { + switch (malformation_) { + case CommitReaderMalformation::DROP_LAST: + reader = std::make_unique(std::move(reader)); + break; + case CommitReaderMalformation::UNSORTED: + reader = std::make_unique(std::move(reader)); + break; + case CommitReaderMalformation::SUBSTITUTE_OFFSET: + reader = std::make_unique(std::move(reader)); + break; + } + } + return readers; + } + Result> AcquireReadView() override { + return delegate_->AcquireReadView(); + } + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override { + return delegate_->CreateQueryReaders(view, offset_begin, context); + } + Status AdvanceCommittedOffset(int64_t committed_offset) override { + return delegate_->AdvanceCommittedOffset(committed_offset); + } + uint64_t GetMemoryUsage() const override { + return delegate_->GetMemoryUsage(); + } + + private: + std::shared_ptr delegate_; + CommitReaderMalformation malformation_; +}; + +class MalformedCoverageRealtimeStoreFactory final : public RealtimeStoreFactory { + public: + explicit MalformedCoverageRealtimeStoreFactory( + CommitReaderMalformation malformation = CommitReaderMalformation::DROP_LAST) + : malformation_(malformation) {} + + Result> Create(RealtimeStoreCreateRequest&& request) override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr delegate, + delegate_.Create(std::move(request))); + return std::shared_ptr( + std::make_shared(delegate, malformation_)); + } + + private: + ArrowRealtimeStoreFactory delegate_; + CommitReaderMalformation malformation_; +}; + +} // namespace namespace { @@ -219,6 +778,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 +813,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 +826,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 +842,7 @@ class RealtimeWriteInteTest : public ::testing::Test { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &c_array)); RecordBatchBuilder builder(&c_array); + builder.SetRowKinds(row_kinds); if (partitioned) { builder.SetPartition({{"pt", partition}}); } @@ -426,6 +1006,16 @@ class RealtimeWriteInteTest : public ::testing::Test { return commit->Expire(); } + Status CommitMessages(const std::vector>& messages, + int64_t commit_identifier) const { + CommitContextBuilder builder(table_path_, commit_user_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, + builder.SetOptions(options_).Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr commit, + FileStoreCommit::Create(std::move(context))); + return commit->Commit(messages, commit_identifier); + } + Result> CreatePlan( const std::shared_ptr& realtime_context, const std::shared_ptr& predicate) const { @@ -546,181 +1136,1322 @@ class RealtimeWriteInteTest : public ::testing::Test { return ReadRows(/*realtime_context=*/nullptr); } - Result CountRows(const std::shared_ptr& plan, - const std::shared_ptr& realtime_context) const { + Result CountRows(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context) const { + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .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))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr count_reader, + table_read->CreateCountReader(plan->Splits())); + return count_reader->CountRows(); + } + + Result GetRealtimeMemoryUsage( + 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()); + uint64_t memory_usage = 0; + for (const RealtimePartitionBucketView& view : views) { + memory_usage += view.store->GetMemoryUsage(); + } + return memory_usage; + } + + Result> ReadPkSequences( + const std::shared_ptr& realtime_context) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, + RealtimeContextImpl::Cast(realtime_context)); + PAIMON_ASSIGN_OR_RAISE(std::vector views, + realtime_context_impl->AcquireReadViews()); + if (views.size() != 1) { + return Status::Invalid("expected one PK real-time read view"); + } + auto read_schema = std::make_unique(); + arrow::FieldVector requested_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())}; + requested_fields.insert(requested_fields.end(), schema_->fields().begin(), + schema_->fields().end()); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema(requested_fields), read_schema.get())); + ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); + RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + PAIMON_ASSIGN_OR_RAISE( + std::vector> readers, + views[0].store->CreateQueryReaders(views[0].read_view, + /*offset_begin=*/0, query_context)); + std::vector sequences; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr imported, + arrow::ImportArray(batch.first.get(), batch.second.get())); + std::shared_ptr values = + std::dynamic_pointer_cast(imported); + if (!values) { + return Status::Invalid("PK query reader did not return a StructArray"); + } + std::shared_ptr sequence_array = + std::dynamic_pointer_cast( + values->GetFieldByName(SpecialFields::SequenceNumber().Name())); + if (!sequence_array) { + return Status::Invalid("PK query reader did not return sequence numbers"); + } + for (int64_t row = 0; row < sequence_array->length(); ++row) { + sequences.push_back(sequence_array->Value(row)); + } + } + reader->Close(); + } + return sequences; + } + + static std::vector> NewFiles( + const std::vector& progresses) { + std::vector> files; + for (const RealtimeCommitProgress& progress : progresses) { + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + if (!message) { + continue; + } + const std::vector>& new_files = + message->GetNewFilesIncrement().NewFiles(); + files.insert(files.end(), new_files.begin(), new_files.end()); + } + return files; + } + + static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { + std::vector seen(static_cast(total_rows), false); + int64_t max_id = -1; + for (const Row& row : rows) { + const auto& [id, payload, partition] = row; + if (id < 0 || id >= total_rows) { + return Status::Invalid("real-time read id is out of range"); + } + if (seen[static_cast(id)]) { + return Status::Invalid("real-time read contains duplicate ids"); + } + if (payload != "value-" + std::to_string(id) || partition != "p0") { + return Status::Invalid("real-time read row does not match its id"); + } + seen[static_cast(id)] = true; + max_id = std::max(max_id, id); + } + for (int64_t id = 0; id <= max_id; ++id) { + if (!seen[static_cast(id)]) { + return Status::Invalid("real-time read contains an id gap"); + } + } + 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_); + PAIMON_ASSIGN_OR_RAISE(std::optional snapshot, snapshot_manager.LatestSnapshot()); + return RealtimeCommitProperties::ReadOffsets(snapshot, options.GetFileSystem()); + } + + void FinalizeCommitAndCheck(FileStoreWrite* writer, + std::vector realtime_commits, + int64_t prepare_identifier, std::vector expected_rows) const { + ASSERT_OK_AND_ASSIGN(std::vector final_commits, + writer->PrepareCommitWithProgress(prepare_identifier)); + realtime_commits.insert(realtime_commits.end(), + std::make_move_iterator(final_commits.begin()), + std::make_move_iterator(final_commits.end())); + ASSERT_OK(Commit(realtime_commits, prepare_identifier)); + ASSERT_OK(writer->Close()); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows()); + 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())}; + schema_ = arrow::schema(fields_); + 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, + CreateRealtimeWriter(realtime_context)); + constexpr int32_t kDate = 19723; + constexpr int64_t kRowCount = 3; + const std::string partition = "2024-01-01"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeDatePartitionBatch(/*first_id=*/0, kRowCount, kDate, partition)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, commits.size()); + ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); + + const std::string normalized_partition = + legacy_partition_name_enabled ? std::to_string(kDate) : partition; + RealtimePartitionBucket partition_bucket({{"pt", normalized_partition}}, /*bucket=*/0); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_before_drop, ReadCommittedOffsets()); + ASSERT_EQ(1, offsets_before_drop.size()); + ASSERT_EQ(kRowCount, offsets_before_drop.at(partition_bucket)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_before_drop, + CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t rows_before_drop, + CountRows(plan_before_drop, /*realtime_context=*/nullptr)); + ASSERT_EQ(kRowCount, rows_before_drop); + + ASSERT_OK(DropPartition({{"pt", partition}}, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_after_drop, ReadCommittedOffsets()); + ASSERT_TRUE(offsets_after_drop.empty()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_after_drop, + CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(int64_t rows_after_drop, + CountRows(plan_after_drop, /*realtime_context=*/nullptr)); + ASSERT_EQ(0, rows_after_drop); + ASSERT_OK(writer->Close()); + } + + void CheckVectorReaderRetry(bool primary_key) { + if (primary_key) { + CreatePkTable(/*partition_keys=*/{"pt"}); + } else { + CreateTable(/*partition_keys=*/{"pt"}); + } + auto close_state = std::make_shared(); + auto factory = std::make_shared(close_state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + 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"); + if (!primary_key) { + ASSERT_EQ(1, close_state->query_close_count->load(std::memory_order_acquire)); + } + + std::vector expected_rows = p0_rows; + expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(expected_rows, actual_rows); + ASSERT_OK(writer->Close()); + } + + std::unique_ptr dir_; + std::string table_path_; + std::string commit_user_ = "realtime_commit_user"; + arrow::FieldVector fields_; + std::shared_ptr schema_; + std::map options_; + std::shared_ptr pool_; +}; + +TEST_F(RealtimeWriteInteTest, TestRealtimeOperationsRequireEnabledOption) { + CreateTable(/*partition_keys=*/{}); + std::map disabled_options = options_; + disabled_options[Options::REALTIME_ENABLED] = "false"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + + WriteContextBuilder write_builder(table_path_, commit_user_); + write_builder.SetOptions(disabled_options) + .WithStreamingMode(true) + .WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, write_builder.Finish()); + ASSERT_NOK_WITH_MSG(FileStoreWrite::Create(std::move(write_context)), + "real-time write requires realtime.enabled=true"); + + ScanContextBuilder scan_builder(table_path_); + scan_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "real-time scan requires realtime.enabled=true"); + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), + "real-time read requires realtime.enabled=true"); + + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(disabled_options).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_NOK_WITH_MSG(commit->CommitWithProgress(/*realtime_commits=*/{}, + /*commit_identifier=*/0, + /*watermark=*/std::nullopt), + "CommitWithProgress requires realtime.enabled=true"); +} + +TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkRead) { + CreatePkTable(); + auto saw_query_predicate = std::make_shared>(false); + auto query_view = std::make_shared>(); + auto factory = + std::make_shared(saw_query_predicate, query_view); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector first_rows = {{1, "old", "p0"}, {2, "two", "p0"}, {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, 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, TestPkMemoryReadAfterSchemaEvolution) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "old", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + std::shared_ptr renamed_payload = arrow::field("renamed_payload", arrow::utf8()); + std::shared_ptr added = arrow::field("added", arrow::int32()); + ASSERT_OK(TestHelper::WriteNextSchema(dir_->GetFileSystem(), table_path_, + {DataField(0, fields_[0]), DataField(1, renamed_payload), + DataField(2, fields_[2]), DataField(3, added)}, + /*highest_field_id=*/3, options_)); + fields_[1] = renamed_payload; + fields_.push_back(added); + schema_ = arrow::schema(fields_); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_OK_AND_ASSIGN(CollectedReadResult result, + ReadPlan(plan, realtime_context, {"id", "renamed_payload", "pt", "added"}, + /*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + ASSERT_EQ(1, result.data->num_chunks()); + std::shared_ptr row = + std::dynamic_pointer_cast(result.data->chunk(0)); + ASSERT_NE(nullptr, row); + ASSERT_EQ(1, row->length()); + std::shared_ptr renamed_values = + std::dynamic_pointer_cast(row->field(2)); + ASSERT_NE(nullptr, renamed_values); + ASSERT_EQ("old", renamed_values->GetString(0)); + ASSERT_TRUE(row->field(4)->IsNull(0)); + result.reader->Close(); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkCompositeMerge) { + CreatePkTable(/*partition_keys=*/{}, /*primary_keys=*/{"id", "payload"}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + 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::make_pair(0, 1)), first_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(0, 2)), first_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t first_snapshot_id, + Commit(first_progress, /*commit_identifier=*/0)); + ASSERT_OK(first_writer->RefreshCommittedSnapshot(first_snapshot_id)); + ASSERT_OK(first_writer->Close()); + first_writer.reset(); + first_context.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr second_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_writer, + CreateRealtimeWriter(second_context)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr p0_second_batch, + MakeBatch({Row{0, "p0-zero-new", "p0"}, Row{2, "p0-two", "p0"}}, + /*partitioned=*/true, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p0_second_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_second_batch, + MakeBatch({Row{10, "p1-ten-deleted", "p1"}, Row{13, "p1-thirteen", "p1"}}, + /*partitioned=*/true, /*bucket=*/1, + {RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(second_writer->Write(std::move(p1_second_batch))); + ASSERT_OK_AND_ASSIGN(std::vector second_progress, + second_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(2, second_progress.size()); + std::map second_ranges; + std::map> second_sequences; + for (const RealtimeCommitProgress& progress : second_progress) { + second_ranges.emplace(progress.partition_bucket, progress.offset_range); + std::shared_ptr message = + std::dynamic_pointer_cast(progress.commit_message); + ASSERT_NE(nullptr, message); + const std::vector>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(1, files.size()); + second_sequences.emplace( + progress.partition_bucket, + std::make_pair(files[0]->min_sequence_number, files[0]->max_sequence_number)); + } + ASSERT_EQ(OffsetRange(2, 4), second_ranges.at(p0b0)); + ASSERT_EQ(OffsetRange(3, 5), second_ranges.at(p1b1)); + ASSERT_EQ((std::make_pair(2, 3)), second_sequences.at(p0b0)); + ASSERT_EQ((std::make_pair(3, 4)), second_sequences.at(p1b1)); + ASSERT_OK_AND_ASSIGN(int64_t second_snapshot_id, + Commit(second_progress, /*commit_identifier=*/1)); + ASSERT_OK(second_writer->RefreshCommittedSnapshot(second_snapshot_id)); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(second_context)); + std::sort(actual_rows.begin(), actual_rows.end()); + ASSERT_EQ((std::vector{{0, "p0-zero-new", "p0"}, + {1, "p0-one", "p0"}, + {2, "p0-two", "p0"}, + {11, "p1-eleven", "p1"}, + {12, "p1-twelve", "p1"}, + {13, "p1-thirteen", "p1"}}), + actual_rows); + ASSERT_OK(second_writer->Close()); + + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(2, offsets.size()); + ASSERT_EQ(4, offsets.at(p0b0)); + ASSERT_EQ(5, offsets.at(p1b1)); +} + +TEST_F(RealtimeWriteInteTest, TestPkRecovery) { + CreatePkTable(); + + 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 failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result> failed_prepare = + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1); + io_hook->Clear(); + ASSERT_TRUE(failed_prepare.status().IsIOError()) << failed_prepare.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failed_prepare, ReadRows()); + ASSERT_EQ((std::vector{{99, "seed", "p0"}}), rows_after_failed_prepare); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr first_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_writer, + CreateRealtimeWriter(first_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(mutations, /*partitioned=*/false, /*bucket=*/0, mutation_kinds)); + ASSERT_OK(first_writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector memory_sequences, ReadPkSequences(first_context)); + ASSERT_EQ((std::vector{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_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); + ASSERT_OK(CommitMessages(compact_messages, /*commit_identifier=*/4)); + ASSERT_OK(compact_writer->Close()); + + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(options_)); + SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); + ASSERT_OK_AND_ASSIGN(std::optional compact_snapshot, + snapshot_manager.LatestSnapshot()); + ASSERT_TRUE(compact_snapshot); + ASSERT_EQ(Snapshot::CommitKind::Compact(), compact_snapshot->GetCommitKind()); + ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets, ReadCommittedOffsets()); + ASSERT_EQ(4, offsets.at(RealtimePartitionBucket(/*partition=*/{}, /*bucket=*/0))); + ASSERT_OK_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, TestPkPluginContract) { + CreatePkTable(); + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{4, "four", "p0"}, Row{3, "three", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch({Row{2, "two", "p0"}, Row{1, "one", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 4), progress[0].offset_range); + ASSERT_OK(Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector rows, ReadRows()); + ASSERT_EQ((std::vector{ + {1, "one", "p0"}, {2, "two", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}), + rows); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsMalformedCoverage) { + CreatePkTable(); + auto factory = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "commit readers did not cover the sealed range"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsEqualCardinalityOffsetSubstitution) { + CreatePkTable(); + auto factory = std::make_shared( + CommitReaderMalformation::SUBSTITUTE_OFFSET); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "duplicate REALTIME_OFFSET"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkRejectsUnsortedPluginRows) { + CreatePkTable(); + auto factory = + std::make_shared(CommitReaderMalformation::UNSORTED); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}, Row{2, "two", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "not globally sorted by primary key and sequence number"); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderClose) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + + auto create_reader = [&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + return table_read->CreateReader(plan->Splits()); + }; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr explicitly_closed_reader, create_reader()); + explicitly_closed_reader->Close(); + explicitly_closed_reader.reset(); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr destroyed_reader, create_reader()); + destroyed_reader.reset(); + ASSERT_EQ(2, state->query_close_count->load(std::memory_order_acquire)); + + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeWriteInteTest, TestPkQueryReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_batch, + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_batch, + MakeBatch({Row{2, "two", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(second_batch))); + + auto create_reader = [&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); ReadContextBuilder read_builder(table_path_); read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) .WithRealtimeContext(realtime_context) .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, TableRead::Create(std::move(read_context))); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr count_reader, - table_read->CreateCountReader(plan->Splits())); - return count_reader->CountRows(); - } - - Result GetRealtimeMemoryUsage( - 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()); - uint64_t memory_usage = 0; - for (const RealtimePartitionBucketView& view : views) { - memory_usage += view.store->GetMemoryUsage(); - } - return memory_usage; - } - - static Status ValidateReadPrefix(const std::vector& rows, int64_t total_rows) { - std::vector seen(static_cast(total_rows), false); - int64_t max_id = -1; - for (const Row& row : rows) { - const auto& [id, payload, partition] = row; - if (id < 0 || id >= total_rows) { - return Status::Invalid("real-time read id is out of range"); - } - if (seen[static_cast(id)]) { - return Status::Invalid("real-time read contains duplicate ids"); - } - if (payload != "value-" + std::to_string(id) || partition != "p0") { - return Status::Invalid("real-time read row does not match its id"); - } - seen[static_cast(id)] = true; - max_id = std::max(max_id, id); - } - for (int64_t id = 0; id <= max_id; ++id) { - if (!seen[static_cast(id)]) { - return Status::Invalid("real-time read contains an id gap"); - } - } - return Status::OK(); - } - - Result ReadCommittedOffsets() const { - PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(options_)); - SnapshotManager snapshot_manager(options.GetFileSystem(), table_path_); - PAIMON_ASSIGN_OR_RAISE(std::optional snapshot, snapshot_manager.LatestSnapshot()); - return RealtimeCommitProperties::ReadOffsets(snapshot, options.GetFileSystem()); - } - - void FinalizeCommitAndCheck(FileStoreWrite* writer, - std::vector realtime_commits, - int64_t prepare_identifier, std::vector expected_rows) const { - ASSERT_OK_AND_ASSIGN(std::vector final_commits, - writer->PrepareCommitWithProgress(prepare_identifier)); - realtime_commits.insert(realtime_commits.end(), - std::make_move_iterator(final_commits.begin()), - std::make_move_iterator(final_commits.end())); - ASSERT_OK(Commit(realtime_commits, prepare_identifier)); - ASSERT_OK(writer->Close()); - - 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())}; - schema_ = arrow::schema(fields_); - 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, - CreateRealtimeWriter(realtime_context)); - constexpr int32_t kDate = 19723; - constexpr int64_t kRowCount = 3; - const std::string partition = "2024-01-01"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeDatePartitionBatch(/*first_id=*/0, kRowCount, kDate, partition)); - ASSERT_OK(writer->Write(std::move(batch))); - ASSERT_OK_AND_ASSIGN(std::vector commits, - writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); - ASSERT_EQ(1, commits.size()); - ASSERT_OK(Commit(commits, /*commit_identifier=*/0)); - - const std::string normalized_partition = - legacy_partition_name_enabled ? std::to_string(kDate) : partition; - RealtimePartitionBucket partition_bucket({{"pt", normalized_partition}}, /*bucket=*/0); - ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_before_drop, ReadCommittedOffsets()); - ASSERT_EQ(1, offsets_before_drop.size()); - ASSERT_EQ(kRowCount, offsets_before_drop.at(partition_bucket)); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_before_drop, - CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); - ASSERT_OK_AND_ASSIGN(int64_t rows_before_drop, - CountRows(plan_before_drop, /*realtime_context=*/nullptr)); - ASSERT_EQ(kRowCount, rows_before_drop); + return table_read->CreateReader(plan->Splits()); + }; - ASSERT_OK(DropPartition({{"pt", partition}}, /*commit_identifier=*/1)); - ASSERT_OK_AND_ASSIGN(RealtimeOffsetMap offsets_after_drop, ReadCommittedOffsets()); - ASSERT_TRUE(offsets_after_drop.empty()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan_after_drop, - CreatePlan(/*realtime_context=*/nullptr, /*predicate=*/nullptr)); - ASSERT_OK_AND_ASSIGN(int64_t rows_after_drop, - CountRows(plan_after_drop, /*realtime_context=*/nullptr)); - ASSERT_EQ(0, rows_after_drop); - ASSERT_OK(writer->Close()); + for (int32_t null_index = 0; null_index <= 1; ++null_index) { + state->query_null_index = null_index; + ASSERT_NOK_WITH_MSG(create_reader(), "PK real-time store returned a null query reader"); + ASSERT_EQ(null_index + 1, state->query_close_count->load(std::memory_order_acquire)); } + ASSERT_OK(writer->Close()); +} - std::unique_ptr dir_; - std::string table_path_; - std::string commit_user_ = "realtime_commit_user"; - arrow::FieldVector fields_; - std::shared_ptr schema_; - std::map options_; - std::shared_ptr pool_; -}; - -TEST_F(RealtimeWriteInteTest, TestRealtimeOperationsRequireEnabledOption) { +TEST_F(RealtimeWriteInteTest, TestAppendQueryReaderCloseFailure) { CreateTable(/*partition_keys=*/{}); - std::map disabled_options = options_; - disabled_options[Options::REALTIME_ENABLED] = "false"; + auto state = std::make_shared(); + state->query_null_index = 1; + auto factory = std::make_shared(state); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - - WriteContextBuilder write_builder(table_path_, commit_user_); - write_builder.SetOptions(disabled_options) - .WithStreamingMode(true) - .WithRealtimeContext(realtime_context); - ASSERT_OK_AND_ASSIGN(std::unique_ptr write_context, write_builder.Finish()); - ASSERT_NOK_WITH_MSG(FileStoreWrite::Create(std::move(write_context)), - "real-time write requires realtime.enabled=true"); - - ScanContextBuilder scan_builder(table_path_); - scan_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); - ASSERT_OK_AND_ASSIGN(std::unique_ptr scan_context, scan_builder.Finish()); - ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), - "real-time scan requires realtime.enabled=true"); + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(disabled_options).WithRealtimeContext(realtime_context); + 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_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), - "real-time read requires realtime.enabled=true"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG(table_read->CreateReader(plan->Splits()), + "append-only real-time store returned a null query reader"); + ASSERT_EQ(1, state->query_close_count->load(std::memory_order_acquire)); - CommitContextBuilder commit_builder(table_path_, commit_user_); - ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, - commit_builder.SetOptions(disabled_options).Finish()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, - FileStoreCommit::Create(std::move(commit_context))); - ASSERT_NOK_WITH_MSG(commit->CommitWithProgress(/*realtime_commits=*/{}, - /*commit_identifier=*/0, - /*watermark=*/std::nullopt), - "CommitWithProgress requires realtime.enabled=true"); + state->query_null_index = -1; + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestAppendCommitAndRead) { - CreateTable(/*partition_keys=*/{}); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, CreateRealtimeWriter()); - std::vector rows = MakeRows(/*first_id=*/0, /*count=*/10, /*partition=*/"p0"); +TEST_F(RealtimeWriteInteTest, TestPkCommitReaderCloseFailure) { + CreatePkTable(); + auto state = std::make_shared(); + state->commit_null_index = 1; + auto factory = std::make_shared(state); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create(factory)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, - MakeBatch(rows, /*partitioned=*/false)); + MakeBatch({Row{1, "one", "p0"}}, /*partitioned=*/false)); ASSERT_OK(writer->Write(std::move(batch))); - FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); + + ASSERT_NOK_WITH_MSG(writer->PrepareCommitWithProgress(/*commit_identifier=*/0), + "PK real-time store returned a null commit reader"); + ASSERT_EQ(1, state->commit_close_count->load(std::memory_order_acquire)); + ASSERT_OK(writer->Close()); } TEST_F(RealtimeWriteInteTest, TestRollingFilesPreserveProgress) { @@ -1154,6 +2885,44 @@ TEST_F(RealtimeWriteInteTest, TestPlanExcludesRowsWrittenAfterMemoryEndOffset) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestPkRejectsReversedVisibleOffsets) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + std::vector rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, /*predicate=*/nullptr)); + ASSERT_EQ(1, plan->Splits().size()); + std::shared_ptr split = + std::dynamic_pointer_cast(plan->Splits()[0]); + ASSERT_NE(nullptr, split); + std::vector> disk_splits = split->DiskSplits(); + std::vector> invalid_splits = {std::make_shared( + split->Version(), split->SnapshotId(), split->Partition(), split->Bucket(), + std::move(disk_splits), split->MemoryEndOffset() + 1, split->MemoryEndOffset(), + split->OpaqueTicket())}; + + ReadContextBuilder read_builder(table_path_); + read_builder.SetOptions(options_) + .SetReadFieldNames({"id", "payload", "pt"}) + .WithRealtimeContext(realtime_context) + .WithMemoryPool(pool_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), + "memory end offset precedes committed end offset"); + + ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); + ASSERT_EQ(rows, actual_rows); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestReadFailsAfterRealtimeSplitTicketExpires) { options_[Options::REALTIME_READ_VIEW_TTL] = "10 ms"; CreateTable(/*partition_keys=*/{}); @@ -1235,50 +3004,12 @@ TEST_F(RealtimeWriteInteTest, TestFailedReaderCreationPreservesRealtimeSplitTick ASSERT_OK(writer->Close()); } -TEST_F(RealtimeWriteInteTest, TestVectorReaderFailurePreservesEarlierSplitTicket) { - CreateTable(/*partition_keys=*/{"pt"}); - ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, - RealtimeContext::Create()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, - CreateRealtimeWriter(realtime_context)); - std::vector p0_rows = MakeRows(/*first_id=*/0, /*count=*/3, /*partition=*/"p0"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p0_batch, - MakeBatch(p0_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p0_batch))); - std::vector p1_rows = MakeRows(/*first_id=*/10, /*count=*/3, /*partition=*/"p1"); - ASSERT_OK_AND_ASSIGN(std::unique_ptr p1_batch, - MakeBatch(p1_rows, /*partitioned=*/true)); - ASSERT_OK(writer->Write(std::move(p1_batch))); - ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, - CreatePlan(realtime_context, /*predicate=*/nullptr)); - ASSERT_EQ(2, plan->Splits().size()); - - std::vector> invalid_splits = plan->Splits(); - std::shared_ptr second_split = - std::dynamic_pointer_cast(invalid_splits[1]); - ASSERT_NE(nullptr, second_split); - std::vector> second_disk_splits = second_split->DiskSplits(); - invalid_splits[1] = std::make_shared( - RealtimeSplit::kCurrentVersion + 1, second_split->SnapshotId(), second_split->Partition(), - second_split->Bucket(), std::move(second_disk_splits), second_split->CommittedEndOffset(), - second_split->MemoryEndOffset(), second_split->OpaqueTicket()); - - ReadContextBuilder read_builder(table_path_); - read_builder.SetOptions(options_) - .SetReadFieldNames({"id", "payload", "pt"}) - .WithRealtimeContext(realtime_context) - .WithMemoryPool(pool_); - ASSERT_OK_AND_ASSIGN(std::unique_ptr read_context, read_builder.Finish()); - ASSERT_OK_AND_ASSIGN(std::unique_ptr table_read, - TableRead::Create(std::move(read_context))); - ASSERT_NOK_WITH_MSG(table_read->CreateReader(invalid_splits), - "unsupported real-time split version"); +TEST_F(RealtimeWriteInteTest, TestAppendVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/false); +} - std::vector expected_rows = p0_rows; - expected_rows.insert(expected_rows.end(), p1_rows.begin(), p1_rows.end()); - ASSERT_OK_AND_ASSIGN(std::vector actual_rows, ReadRows(plan, realtime_context)); - ASSERT_EQ(expected_rows, actual_rows); - ASSERT_OK(writer->Close()); +TEST_F(RealtimeWriteInteTest, TestPkVectorRetry) { + CheckVectorReaderRetry(/*primary_key=*/true); } TEST_F(RealtimeWriteInteTest, TestCloseWriterKeepsContextReadable) { @@ -2170,8 +3901,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, @@ -2407,6 +4142,14 @@ TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { ASSERT_OK(writer->Close()); } +TEST_F(RealtimeWriteInteTest, TestConcurrentWritePrepareCommitReadAndRefresh) { + RunConcurrencyTest(/*primary_key=*/false); +} + +TEST_F(RealtimeWriteInteTest, TestPkConcurrency) { + RunConcurrencyTest(/*primary_key=*/true); +} + TEST_F(RealtimeWriteInteTest, TestMultiplePartitions) { CreateTable(/*partition_keys=*/{"pt"}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -2711,4 +4454,100 @@ TEST_F(RealtimeWriteInteTest, TestRestoreOffsetFromCommittedSnapshot) { ASSERT_EQ(5, second_committed_offsets.at(partition_bucket)); } +TEST_F(RealtimeWriteInteTest, TestPkExternalCommitRecovery) { + CreatePkTable(); + const std::vector seed_rows = {{99, "seed", "p0"}}; + ReplayPkWalAndCommit(seed_rows, /*row_kinds=*/{}, /*commit_identifier=*/0, seed_rows); + + const std::vector wal = { + {1, "one", "p0"}, {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector row_kinds = { + RecordBatch::RowKind::INSERT, RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_batch, + MakeBatch(wal, /*partitioned=*/false, /*bucket=*/0, row_kinds)); + ASSERT_OK(failed_writer->Write(std::move(failed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector failed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, failed_progress.size()); + CommitContextBuilder commit_builder(table_path_, commit_user_); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit_context, + commit_builder.SetOptions(options_).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr commit, + FileStoreCommit::Create(std::move(commit_context))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Result failed_commit = + commit->CommitWithProgress(failed_progress, /*commit_identifier=*/1, + /*watermark=*/std::nullopt); + io_hook->Clear(); + ASSERT_TRUE(failed_commit.status().IsIOError()) << failed_commit.status().ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(seed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {99, "seed", "p0"}}; + ReplayPkWalAndCommit(wal, row_kinds, /*commit_identifier=*/1, expected_rows); +} + +TEST_F(RealtimeWriteInteTest, TestPkRefreshRecovery) { + CreatePkTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr failed_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr failed_writer, + CreateRealtimeWriter(failed_context)); + + const std::vector base_rows = {{1, "one", "p0"}, {2, "two", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr base_batch, + MakeBatch(base_rows, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(base_batch))); + ASSERT_OK_AND_ASSIGN(std::vector base_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t base_snapshot_id, Commit(base_progress, /*commit_identifier=*/0)); + ASSERT_OK(failed_writer->RefreshCommittedSnapshot(base_snapshot_id)); + + const std::vector committed_wal = { + {1, "one-new", "p0"}, {2, "deleted", "p0"}, {3, "three", "p0"}}; + const std::vector committed_kinds = {RecordBatch::RowKind::UPDATE_AFTER, + RecordBatch::RowKind::DELETE, + RecordBatch::RowKind::INSERT}; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr committed_batch, + MakeBatch(committed_wal, /*partitioned=*/false, /*bucket=*/0, committed_kinds)); + ASSERT_OK(failed_writer->Write(std::move(committed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector committed_progress, + failed_writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t committed_snapshot_id, + Commit(committed_progress, /*commit_identifier=*/1)); + + const std::vector replay_wal = {{4, "four", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr replay_batch, + MakeBatch(replay_wal, /*partitioned=*/false)); + ASSERT_OK(failed_writer->Write(std::move(replay_batch))); + IOHook* io_hook = IOHook::GetInstance(); + ScopeGuard hook_guard([io_hook]() { io_hook->Clear(); }); + io_hook->Reset(/*pos=*/0, IOHook::Mode::RETURN_ERROR); + Status failed_refresh = failed_writer->RefreshCommittedSnapshot(committed_snapshot_id); + io_hook->Clear(); + ASSERT_TRUE(failed_refresh.IsIOError()) << failed_refresh.ToString(); + ASSERT_OK(failed_writer->Close()); + failed_writer.reset(); + failed_context.reset(); + const std::vector committed_rows = {{1, "one-new", "p0"}, {3, "three", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector rows_after_failure, ReadRows()); + ASSERT_EQ(committed_rows, rows_after_failure); + + const std::vector expected_rows = { + {1, "one-new", "p0"}, {3, "three", "p0"}, {4, "four", "p0"}}; + ReplayPkWalAndCommit(replay_wal, /*row_kinds=*/{}, /*commit_identifier=*/2, expected_rows); +} + } // namespace paimon::test