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