lxy-9602 commented on code in PR #224:
URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3850564459
##########
src/paimon/core/operation/file_store_write.cpp:
##########
@@ -197,7 +198,26 @@ 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(ValidatePrimaryKeyRealtimeOptions(options,
*schema));
+ if (ignore_previous_files) {
+ return Status::NotImplemented(
+ "PK realtime v1 requires restore from the latest
snapshot");
+ }
+ if (!ctx->GetWriteSchema().empty()) {
+ return Status::NotImplemented(
+ "PK realtime v1 does not support a custom write schema");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::optional<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));
+ }
Review Comment:
Could some of the checks here be extracted into a shared func and reused
with `GetRealtimeContext` for append tables? It seems like
`AdvanceCommittedProgress` may also contain similar logic.
##########
src/paimon/core/table/source/key_value_table_read.cpp:
##########
@@ -34,6 +50,62 @@ class Executor;
class FileStorePathFactory;
class InternalReadContext;
class MemoryPool;
+struct ColumnarBatchContext;
+
+namespace {
+
+Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> CreateMemoryReaders(
+ const std::shared_ptr<RealtimeSplit>& split, const
RealtimePartitionBucketView& memory,
+ 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) {
+ std::shared_ptr<arrow::Schema> full_value_schema =
+
DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields());
+ arrow::FieldVector prepared_fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())
+ ->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())};
+ prepared_fields.insert(prepared_fields.end(),
full_value_schema->fields().begin(),
+ full_value_schema->fields().end());
+ std::shared_ptr<arrow::Schema> prepared_schema =
arrow::schema(std::move(prepared_fields));
+ auto c_schema = std::make_unique<ArrowSchema>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema,
c_schema.get()));
+ ScopeGuard schema_guard([schema = c_schema.get()]() {
ArrowSchemaRelease(schema); });
+ RealtimeQueryContext query_context{c_schema.get(), nullptr, false};
+ PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<BatchReader>>
batch_readers,
+ memory.store->CreateQueryReaders(memory.read_view,
0, query_context));
+ ScopeGuard batch_readers_guard([&batch_readers]() {
+ for (const std::unique_ptr<BatchReader>& reader : batch_readers) {
+ if (reader) {
+ reader->Close();
+ }
+ }
+ });
+ std::vector<std::unique_ptr<KeyValueRecordReader>> result;
+ result.reserve(batch_readers.size());
+ for (std::unique_ptr<BatchReader>& reader : batch_readers) {
+ if (!reader) {
+ return Status::Invalid("PK real-time store returned a null query
reader");
+ }
+ PAIMON_ASSIGN_OR_RAISE(
+ std::unique_ptr<KeyValueRecordReader> prepared_reader,
+ AdaptPreparedBatchReader(
+ std::move(reader), prepared_schema,
+ OffsetRange(split->CommittedEndOffset(),
split->MemoryEndOffset()), key_schema,
+ value_schema, key_comparator, memory_pool));
+ auto merge = std::make_unique<DeduplicateMergeFunction>(false);
+ result.push_back(std::make_unique<MergedKeyValueRecordReader>(
+ std::move(prepared_reader), key_comparator,
Review Comment:
Could we use `CreateMergeFunction` from `primary_key_table_utils.h` to
create the merge function here?
##########
include/paimon/realtime/realtime_store.h:
##########
@@ -41,10 +43,42 @@ namespace paimon {
class MemoryPool;
class Predicate;
-/// A table record batch and its framework-assigned contiguous offset range.
+struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig {
+ StatisticsMode statistics_mode;
+};
+
+struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig {
+ /// Primary-key fields after removing partition fields, in comparison
order.
+ std::vector<std::string> trimmed_primary_keys;
+};
+
+using RealtimeStoreCreateConfig =
+ std::variant<AppendRealtimeStoreCreateConfig,
PrimaryKeyRealtimeStoreCreateConfig>;
+
+/// 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 prepared 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;
+ /// Partition values identifying the store.
+ std::map<std::string, std::string> partition;
+ /// Bucket identifying the store within its partition.
+ int32_t bucket = -1;
+ /// Mode-specific store configuration.
+ RealtimeStoreCreateConfig mode_config;
Review Comment:
Does the store currently use `partition` or `bucket` internally? These
values seem to belong to the framework-side store registry and
committed-progress identity rather than the store plugin contract. I suggest
changing `GetOrCreateRealtimeStore` to accept a `RealtimeStoreCreateRequest`
and a separate `RealtimePartitionBucket`, and removing `partition` and `bucket`
from `RealtimeStoreCreateRequest`.
##########
src/paimon/core/table/source/key_value_table_read.cpp:
##########
@@ -34,6 +50,62 @@ class Executor;
class FileStorePathFactory;
class InternalReadContext;
class MemoryPool;
+struct ColumnarBatchContext;
+
+namespace {
+
+Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> CreateMemoryReaders(
+ const std::shared_ptr<RealtimeSplit>& split, const
RealtimePartitionBucketView& memory,
+ 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) {
+ std::shared_ptr<arrow::Schema> full_value_schema =
+
DataField::ConvertDataFieldsToArrowSchema(context->GetTableSchema()->Fields());
+ arrow::FieldVector prepared_fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())
+ ->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())};
+ prepared_fields.insert(prepared_fields.end(),
full_value_schema->fields().begin(),
+ full_value_schema->fields().end());
+ std::shared_ptr<arrow::Schema> prepared_schema =
arrow::schema(std::move(prepared_fields));
+ auto c_schema = std::make_unique<ArrowSchema>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*prepared_schema,
c_schema.get()));
+ ScopeGuard schema_guard([schema = c_schema.get()]() {
ArrowSchemaRelease(schema); });
+ RealtimeQueryContext query_context{c_schema.get(), nullptr, false};
+ PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<BatchReader>>
batch_readers,
+ memory.store->CreateQueryReaders(memory.read_view,
0, query_context));
Review Comment:
Do we really need to fetch all fields here? In theory, wouldn’t it be enough
to populate something like `MergeFileSplitRead` with the PK fields plus the
sequence field?
##########
src/paimon/core/schema/schema_validation_test.cpp:
##########
@@ -46,6 +46,13 @@ TEST(SchemaValidationTest, TestSimple) {
ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema));
}
+TEST(SchemaValidationTest, TestRealtimeOffsetIsNotGloballyReserved) {
+ auto schema = arrow::schema({arrow::field("_REALTIME_OFFSET",
arrow::int64())});
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableSchema> table_schema,
+ TableSchema::Create(0, schema, {}, {}, {}));
+ ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema));
+}
Review Comment:
I’d suggest keeping `_REALTIME_OFFSET` as a system field for now. From the
current implementation, it seems writes are validated based on the realtime
context, and I’m concerned this assumption might be overlooked in future
development. If we eventually do run into a real business field using this
name, we can revisit and adjust then.
##########
src/paimon/core/operation/merge_file_split_read.cpp:
##########
@@ -78,6 +78,126 @@ struct KeyValue;
template <typename T>
class MergeFunctionWrapper;
+class MergeFileSplitRead::RealtimeReaderBuilder {
+ public:
+ static Result<std::unique_ptr<BatchReader>> Create(
+ MergeFileSplitRead* owner, const std::vector<std::shared_ptr<Split>>&
disk_splits,
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&&
additional_readers) {
Review Comment:
Could we move the output parameter `MergeFileSplitRead* owner` to the end of
the parameter list?
##########
src/paimon/core/realtime/primary_key_realtime_store.cpp:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 <mutex>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "paimon/common/metrics/metrics_impl.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/core_options.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/schema/table_schema.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const
TableSchema& schema) {
+ if (options.GetBucket() <= 0) {
+ return Status::NotImplemented("PK realtime v1 requires fixed buckets");
+ }
+ if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) {
+ return Status::NotImplemented("PK realtime v1 supports only the
DEDUPLICATE merge engine");
+ }
+ if (options.DataEvolutionEnabled()) {
+ return Status::NotImplemented("PK realtime v1 does not support data
evolution");
+ }
+ if (!options.GetFieldsSequenceGroups().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence groups");
+ }
+ if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete()
||
+ options.AggregationRemoveRecordOnDelete() ||
+ !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) {
+ return Status::NotImplemented("PK realtime v1 requires default delete
behavior");
+ }
+ if (!options.GetSequenceField().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence.field");
+ }
+ if (!options.SequenceFieldSortOrderIsAscending()) {
+ return Status::NotImplemented(
+ "PK realtime v1 supports only ascending
sequence.field.sort-order");
+ }
+ if (options.NeedLookup() || options.DeletionVectorsEnabled() ||
+ options.GetChangelogProducer() != ChangelogProducer::NONE) {
+ return Status::NotImplemented("PK realtime v1 does not support lookup
or early MOR");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::vector<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 v1 does not support FLOAT or DOUBLE primary
keys");
+ }
+ }
+ if (options.GlobalIndexEnabled()) {
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions,
+ PrimaryKeyIndexDefinitions::Create(schema));
+ if (!definitions.Definitions().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
global indexes");
+ }
+ }
+ return Status::OK();
+}
+
+namespace {
+
+uint64_t GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data) {
+ uint64_t total = 0;
Review Comment:
`GetArrayMemoryUsage` is used by both append and PK paths. Could we move it
into `arrow_utils` so it can be shared?
##########
src/paimon/core/realtime/prepared_key_value_reader.cpp:
##########
@@ -0,0 +1,738 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "arrow/array/array_base.h"
+#include "arrow/array/array_nested.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/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/utils/nested_projection_utils.h"
+#include "paimon/macros.h"
+#include "paimon/reader/batch_reader.h"
+#include "paimon/status.h"
+
+namespace paimon {
+
+namespace {
+
+constexpr int32_t kValueKindIndex = 0;
+constexpr int32_t kSequenceNumberIndex = 1;
+constexpr int32_t kRealtimeOffsetIndex = 2;
+constexpr int32_t kPreparedValueStartIndex = 3;
+
+template <typename Reader>
+void CloseReaders(const std::vector<std::unique_ptr<Reader>>& readers) {
+ for (const std::unique_ptr<Reader>& reader : readers) {
+ if (reader) {
+ reader->Close();
+ }
+ }
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignArrayByPaimonIds(
+ const std::shared_ptr<arrow::Array>& array, const
std::shared_ptr<arrow::DataType>& read_type,
+ arrow::MemoryPool* arrow_pool);
+
+class RealtimeOffsetCoverage {
+ public:
+ static Result<std::shared_ptr<RealtimeOffsetCoverage>> Create(const
OffsetRange& sealed_offsets,
+ size_t
reader_count) {
+ if (sealed_offsets.begin < 0 || sealed_offsets.end <
sealed_offsets.begin) {
+ return Status::Invalid("PK real-time store returned an invalid
sealed offset range");
+ }
+ return std::shared_ptr<RealtimeOffsetCoverage>(
+ new RealtimeOffsetCoverage(sealed_offsets, reader_count));
+ }
+
+ Status Add(const arrow::Int64Array& offsets) {
+ for (int64_t row = 0; row < offsets.length(); ++row) {
+ const int64_t offset = offsets.Value(row);
+ if (offset < sealed_offsets_.begin || offset >=
sealed_offsets_.end) {
+ return Status::Invalid(
+ "PK real-time store commit reader offset is outside the
sealed range");
+ }
+ min_seen_offset_ = std::min(min_seen_offset_, offset);
+ max_seen_offset_ = std::max(max_seen_offset_, offset);
+ ++seen_count_;
+ }
+ return Status::OK();
+ }
+
+ Status FinishReader() {
+ ++finished_reader_count_;
+ if (finished_reader_count_ == reader_count_ &&
+ (seen_count_ != sealed_offsets_.Count() ||
+ (seen_count_ > 0 && (min_seen_offset_ != sealed_offsets_.begin ||
+ max_seen_offset_ != sealed_offsets_.end -
1)))) {
+ return Status::Invalid(
+ "PK real-time store commit readers did not cover the sealed
range");
+ }
+ return Status::OK();
+ }
+
+ private:
+ RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t
reader_count)
+ : sealed_offsets_(sealed_offsets), reader_count_(reader_count) {}
+
+ OffsetRange sealed_offsets_;
+ size_t reader_count_;
+ int64_t min_seen_offset_ = std::numeric_limits<int64_t>::max();
+ int64_t max_seen_offset_ = std::numeric_limits<int64_t>::min();
+ int64_t seen_count_ = 0;
+ size_t finished_reader_count_ = 0;
+};
+
+Status CheckPreparedField(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("prepared schema missing transport
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(
+ "prepared schema field {} must be non-null {}:{} with field id {},
got {}:{} "
+ "nullable={} field id {}",
+ field_idx, expected_field.Name(),
expected_field.Type()->ToString(),
+ expected_field.Id(), field->name(), field->type()->ToString(),
field->nullable(),
+ field_id));
+ }
+ return Status::OK();
+}
+
+Result<int32_t> FindFieldIndexByPaimonId(const arrow::FieldVector& fields,
int32_t field_id) {
+ std::optional<int32_t> matching_index;
+ for (int32_t i = 0; i < static_cast<int32_t>(fields.size()); ++i) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t candidate_id,
+
NestedProjectionUtils::GetPaimonFieldId(fields[i]));
+ if (candidate_id == field_id) {
+ if (matching_index.has_value()) {
+ return Status::Invalid(
+ fmt::format("duplicate field id {} in prepared schema",
field_id));
+ }
+ matching_index = i;
+ }
+ }
+ if (matching_index.has_value()) {
+ return matching_index.value();
+ }
+ return Status::Invalid(fmt::format("cannot find field id {} in prepared
schema", field_id));
+}
+
+Status ValidateProjectionType(const std::shared_ptr<arrow::DataType>&
prepared_type,
+ const std::shared_ptr<arrow::DataType>&
query_type) {
+ if (prepared_type->id() != query_type->id()) {
+ return Status::Invalid(fmt::format("prepared value type {} does not
match query type {}",
+ prepared_type->ToString(),
query_type->ToString()));
+ }
+ switch (query_type->id()) {
+ case arrow::Type::STRUCT: {
+ const arrow::FieldVector& prepared_fields =
prepared_type->fields();
+ for (const std::shared_ptr<arrow::Field>& query_field :
query_type->fields()) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t query_id,
+
NestedProjectionUtils::GetPaimonFieldId(query_field));
+ PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx,
+
FindFieldIndexByPaimonId(prepared_fields, query_id));
+
PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_fields[prepared_idx]->type(),
+
query_field->type()));
+ }
+ return Status::OK();
+ }
+ case arrow::Type::LIST:
+ return ValidateProjectionType(prepared_type->field(0)->type(),
+ query_type->field(0)->type());
+ case arrow::Type::MAP: {
+ const std::shared_ptr<arrow::MapType> prepared_map =
+ checked_pointer_cast<arrow::MapType>(prepared_type);
+ const std::shared_ptr<arrow::MapType> query_map =
+ checked_pointer_cast<arrow::MapType>(query_type);
+ PAIMON_RETURN_NOT_OK(
+ ValidateProjectionType(prepared_map->key_type(),
query_map->key_type()));
+ return ValidateProjectionType(prepared_map->item_type(),
query_map->item_type());
+ }
+ default:
+ if (!prepared_type->Equals(*query_type)) {
+ return Status::Invalid(
+ fmt::format("prepared leaf type {} does not match query
type {}",
+ prepared_type->ToString(),
query_type->ToString()));
+ }
+ return Status::OK();
+ }
+}
+
+Status ValidateProjectionSchema(const std::shared_ptr<arrow::Schema>&
prepared_schema,
+ const std::shared_ptr<arrow::Schema>&
query_schema) {
+ arrow::FieldVector prepared_value_fields(
+ prepared_schema->fields().begin() + kPreparedValueStartIndex,
+ prepared_schema->fields().end());
+ for (const std::shared_ptr<arrow::Field>& query_field :
query_schema->fields()) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t query_id,
+
NestedProjectionUtils::GetPaimonFieldId(query_field));
+ PAIMON_ASSIGN_OR_RAISE(int32_t prepared_idx,
+ FindFieldIndexByPaimonId(prepared_value_fields,
query_id));
+
PAIMON_RETURN_NOT_OK(ValidateProjectionType(prepared_value_fields[prepared_idx]->type(),
+ query_field->type()));
+ }
+ return Status::OK();
+}
+
+Status ValidateExactCommitSchema(const std::shared_ptr<arrow::Schema>&
prepared_schema,
+ const std::shared_ptr<arrow::Schema>&
value_schema) {
+ if (prepared_schema->num_fields() != value_schema->num_fields() +
kPreparedValueStartIndex) {
+ return Status::Invalid("commit requires the exact prepared writer
schema");
+ }
+ for (int32_t i = 0; i < value_schema->num_fields(); ++i) {
+ if (!prepared_schema->field(i + kPreparedValueStartIndex)
+ ->Equals(value_schema->field(i), true)) {
+ return Status::Invalid("commit requires the exact prepared writer
schema");
+ }
+ }
+ return Status::OK();
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignStructArrayByPaimonIds(
+ const std::shared_ptr<arrow::StructArray>& array,
+ const std::shared_ptr<arrow::StructType>& read_type, arrow::MemoryPool*
arrow_pool) {
+ const std::shared_ptr<arrow::StructType> data_type =
+ checked_pointer_cast<arrow::StructType>(array->type());
+ std::unordered_map<int32_t, int32_t> data_field_id_to_idx;
+ data_field_id_to_idx.reserve(data_type->num_fields());
+ for (int32_t i = 0; i < data_type->num_fields(); ++i) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
+
NestedProjectionUtils::GetPaimonFieldId(data_type->field(i)));
+ if (!data_field_id_to_idx.emplace(field_id, i).second) {
+ return Status::Invalid(
+ fmt::format("duplicate field id {} in prepared value struct",
field_id));
+ }
+ }
+
+ arrow::ArrayVector aligned_arrays;
+ aligned_arrays.reserve(read_type->num_fields());
+ for (const std::shared_ptr<arrow::Field>& read_field :
read_type->fields()) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t read_field_id,
+
NestedProjectionUtils::GetPaimonFieldId(read_field));
+ auto data_iter = data_field_id_to_idx.find(read_field_id);
+ if (data_iter == data_field_id_to_idx.end()) {
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Array> null_child,
+ arrow::MakeArrayOfNull(read_field->type(), array->offset() +
array->length(),
+ arrow_pool));
+ aligned_arrays.push_back(std::move(null_child));
+ continue;
+ }
+ std::shared_ptr<arrow::Array> child =
+ arrow::MakeArray(array->data()->child_data[data_iter->second]);
+ PAIMON_ASSIGN_OR_RAISE(child, AlignArrayByPaimonIds(child,
read_field->type(), arrow_pool));
+ aligned_arrays.push_back(std::move(child));
+ }
+
+ std::shared_ptr<arrow::ArrayData> aligned_data = array->data()->Copy();
+ aligned_data->type = read_type;
+ aligned_data->child_data.clear();
+ aligned_data->child_data.reserve(aligned_arrays.size());
+ for (const std::shared_ptr<arrow::Array>& aligned_array : aligned_arrays) {
+ aligned_data->child_data.push_back(aligned_array->data());
+ }
+ return arrow::MakeArray(std::move(aligned_data));
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignListArrayByPaimonIds(
+ const std::shared_ptr<arrow::ListArray>& array,
+ const std::shared_ptr<arrow::ListType>& read_type, arrow::MemoryPool*
arrow_pool) {
+ std::shared_ptr<arrow::Array> values = array->values();
+ PAIMON_ASSIGN_OR_RAISE(values,
+ AlignArrayByPaimonIds(values,
read_type->value_type(), arrow_pool));
+ std::shared_ptr<arrow::ArrayData> new_data = array->data()->Copy();
+ new_data->type = read_type;
+ new_data->child_data = {values->data()};
+ return arrow::MakeArray(new_data);
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignMapArrayByPaimonIds(
+ const std::shared_ptr<arrow::MapArray>& array, const
std::shared_ptr<arrow::MapType>& read_type,
+ arrow::MemoryPool* arrow_pool) {
+ std::shared_ptr<arrow::Array> keys = array->keys();
+ PAIMON_ASSIGN_OR_RAISE(keys, AlignArrayByPaimonIds(keys,
read_type->key_type(), arrow_pool));
+ std::shared_ptr<arrow::Array> items = array->items();
+ PAIMON_ASSIGN_OR_RAISE(items, AlignArrayByPaimonIds(items,
read_type->item_type(), arrow_pool));
+
+ const std::shared_ptr<arrow::ArrayData>& entries_data =
array->data()->child_data[0];
+ std::shared_ptr<arrow::ArrayData> new_entries = entries_data->Copy();
+ new_entries->type = arrow::struct_({read_type->key_field(),
read_type->item_field()});
+ new_entries->child_data = {keys->data(), items->data()};
+
+ std::shared_ptr<arrow::ArrayData> new_data = array->data()->Copy();
+ new_data->type = read_type;
+ new_data->child_data = {std::move(new_entries)};
+ return arrow::MakeArray(new_data);
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignArrayByPaimonIds(
+ const std::shared_ptr<arrow::Array>& array, const
std::shared_ptr<arrow::DataType>& read_type,
+ arrow::MemoryPool* arrow_pool) {
+ if (array->type()->id() != read_type->id()) {
+ return Status::Invalid(fmt::format("prepared value type {} does not
match query type {}",
+ array->type()->ToString(),
read_type->ToString()));
+ }
+ switch (read_type->id()) {
+ case arrow::Type::STRUCT:
+ return
AlignStructArrayByPaimonIds(checked_pointer_cast<arrow::StructArray>(array),
+
checked_pointer_cast<arrow::StructType>(read_type),
+ arrow_pool);
+ case arrow::Type::LIST:
+ return
AlignListArrayByPaimonIds(checked_pointer_cast<arrow::ListArray>(array),
+
checked_pointer_cast<arrow::ListType>(read_type),
+ arrow_pool);
+ case arrow::Type::MAP:
+ return
AlignMapArrayByPaimonIds(checked_pointer_cast<arrow::MapArray>(array),
+
checked_pointer_cast<arrow::MapType>(read_type),
+ arrow_pool);
+ default:
+ if (!array->type()->Equals(*read_type)) {
+ return Status::Invalid(
+ fmt::format("prepared leaf type {} does not match query
type {}",
+ array->type()->ToString(),
read_type->ToString()));
+ }
+ return array;
+ }
+}
+
+Result<arrow::ArrayVector> ProjectFieldsByPaimonIds(
+ const std::shared_ptr<arrow::StructArray>& data_batch,
+ const std::shared_ptr<arrow::Schema>& prepared_schema,
+ const std::shared_ptr<arrow::Schema>& query_schema, arrow::MemoryPool*
arrow_pool) {
+ std::unordered_map<int32_t, int32_t> prepared_field_id_to_idx;
+ prepared_field_id_to_idx.reserve(prepared_schema->num_fields());
+ for (int32_t i = kPreparedValueStartIndex; i <
prepared_schema->num_fields(); ++i) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
+
NestedProjectionUtils::GetPaimonFieldId(prepared_schema->field(i)));
+ if (!prepared_field_id_to_idx.emplace(field_id, i).second) {
+ return Status::Invalid(
+ fmt::format("duplicate field id {} in prepared schema",
field_id));
+ }
+ }
+
+ arrow::ArrayVector result;
+ result.reserve(query_schema->num_fields());
+ for (const std::shared_ptr<arrow::Field>& query_field :
query_schema->fields()) {
+ PAIMON_ASSIGN_OR_RAISE(int32_t query_field_id,
+
NestedProjectionUtils::GetPaimonFieldId(query_field));
+ auto prepared_iter = prepared_field_id_to_idx.find(query_field_id);
+ if (prepared_iter == prepared_field_id_to_idx.end()) {
+ return Status::Invalid(
+ fmt::format("cannot find field id {} in prepared schema",
query_field_id));
+ }
+ std::shared_ptr<arrow::Array> field_array =
data_batch->field(prepared_iter->second);
+ PAIMON_ASSIGN_OR_RAISE(field_array,
+ AlignArrayByPaimonIds(field_array,
query_field->type(), arrow_pool));
+ result.push_back(std::move(field_array));
+ }
+ return result;
+}
+
+class PreparedKeyValueReader final : public KeyValueRecordReader {
+ public:
+ PreparedKeyValueReader(std::unique_ptr<BatchReader>&& reader,
+ const std::shared_ptr<arrow::Schema>&
prepared_schema,
+ const std::optional<OffsetRange>& visible_offsets,
+ 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<MemoryPool>& pool,
+ const std::shared_ptr<RealtimeOffsetCoverage>&
offset_coverage)
+ : reader_(std::move(reader)),
+ prepared_schema_(prepared_schema),
+ visible_offsets_(visible_offsets),
+ key_schema_(key_schema),
+ value_schema_(value_schema),
+ key_comparator_(key_comparator),
+ pool_(pool),
+ arrow_pool_(GetArrowPool(pool)),
+ offset_coverage_(offset_coverage) {}
+
+ ~PreparedKeyValueReader() override {
+ Close();
+ }
+
+ class Iterator final : public KeyValueRecordReader::Iterator {
+ public:
+ explicit Iterator(PreparedKeyValueReader* reader) : reader_(reader) {}
+
+ Result<bool> HasNext() const override {
+ return cursor_ < reader_->RowCount();
+ }
+
+ Result<KeyValue> Next() override {
+ if (cursor_ >= reader_->RowCount()) {
+ return Status::Invalid("No more prepared 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:
+ PreparedKeyValueReader* reader_;
+ int64_t cursor_ = 0;
+ };
+
+ Result<std::unique_ptr<KeyValueRecordReader::Iterator>> NextBatch()
override {
+ if (first_error_.has_value()) {
+ return first_error_.value();
+ }
+ Result<std::unique_ptr<KeyValueRecordReader::Iterator>> result =
NextBatchImpl();
+ if (!result.ok()) {
+ first_error_ = result.status();
+ Close();
+ }
+ return result;
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return reader_->GetReaderMetrics();
+ }
+
+ void Close() override {
+ if (closed_) {
+ return;
+ }
+ closed_ = true;
+ ResetBatchState();
+ reader_->Close();
+ }
+
+ private:
+ Result<std::unique_ptr<KeyValueRecordReader::Iterator>> NextBatchImpl() {
+ while (true) {
+ ResetBatchState();
+ PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch,
reader_->NextBatch());
+ if (BatchReader::IsEofBatch(batch)) {
+ if (offset_coverage_ && !offset_coverage_finished_) {
+ offset_coverage_finished_ = true;
+ PAIMON_RETURN_NOT_OK(offset_coverage_->FinishReader());
+ }
+ return std::unique_ptr<KeyValueRecordReader::Iterator>();
+ }
+ 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 prepared batch to
StructArray");
+ }
+ std::shared_ptr<arrow::StructArray> data_batch =
+ checked_pointer_cast<arrow::StructArray>(arrow_array);
+ Status transport_status =
+
ValidatePreparedTransportSchema(arrow::schema(data_batch->type()->fields()));
+ if (!transport_status.ok()) {
+ return Status::Invalid(
+ "prepared batch field does not match prepared transport "
+ "schema: ",
+ transport_status.ToString());
+ }
+ if (visible_offsets_.has_value()) {
+ PAIMON_RETURN_NOT_OK(ValidateProjectionSchema(
Review Comment:
Could we simplify this to a normal projection using precomputed field
indexes?
The real-time path could assume that the in-memory store uses the schema
associated with the current real-time/read context. If the table schema
changes, the caller should recreate the real-time store/context and rebuild the
read context instead of attempting schema-evolution reconciliation here.
Nested-field projection should also be handled by the store plugin according
to `RealtimeQueryContext::read_schema`. The plugin should return batches
matching the requested prepared schema exactly.
With that contract, `PreparedKeyValueReader` only needs to select key and
value columns and construct `KeyValue` records.
The field-ID-based lookup, schema compatibility validation, missing-field
null filling, nested-column reshaping, and type reconciliation appear
unnecessary here and make the adapter significantly more complicated.
##########
src/paimon/core/realtime/realtime_primary_key_writer.cpp:
##########
@@ -0,0 +1,325 @@
+/*
+ * 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/data_field.h"
+#include "paimon/common/types/row_kind.h"
+#include "paimon/common/utils/arrow/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
+#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
+#include "paimon/core/mergetree/merge_tree_writer.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/realtime/realtime_context_impl.h"
+#include "paimon/core/utils/commit_increment.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+namespace {
+
+struct PreparedArrayPrivateData {
+ void (*release)(ArrowArray*);
+ void* private_data;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool;
+};
+
+void ReleasePreparedArray(ArrowArray* array) {
+ auto* data = static_cast<PreparedArrayPrivateData*>(array->private_data);
+ array->release = data->release;
+ array->private_data = data->private_data;
+ array->release(array);
+ delete data;
+}
+
+Status RetainPreparedArrayPool(ArrowArray* array,
+ const std::shared_ptr<arrow::MemoryPool>&
arrow_pool) {
+ if (!array || !array->release || !arrow_pool) {
+ return Status::Invalid("cannot retain prepared batch memory pool");
+ }
+ array->private_data =
+ new PreparedArrayPrivateData{array->release, array->private_data,
arrow_pool};
+ array->release = ReleasePreparedArray;
+ return Status::OK();
+}
+
+Result<std::shared_ptr<arrow::StructArray>> PrepareBatch(
+ std::unique_ptr<RecordBatch>&& batch, const
std::shared_ptr<arrow::Schema>& write_schema,
+ const std::shared_ptr<arrow::Schema>& prepared_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> prepared,
+ arrow::StructArray::Make(std::move(columns),
prepared_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(prepared), options,
&context));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum sorted,
+ arrow::compute::Take(arrow::Datum(prepared), indices,
+ arrow::compute::TakeOptions::NoBoundsCheck(),
&context));
+ std::shared_ptr<arrow::Array> sorted_array = sorted.make_array();
+ if (!sorted_array || sorted_array->type_id() != arrow::Type::STRUCT) {
+ return Status::Invalid("PK real-time sorted batch is not a
StructArray");
+ }
+ return checked_pointer_cast<arrow::StructArray>(std::move(sorted_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::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ 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 (!store_state.store || !merge_tree_writer || !write_schema ||
!key_comparator ||
+ !realtime_context || !memory_pool) {
+ return Status::Invalid("PK real-time writer received a null
dependency");
+ }
+ if (trimmed_primary_keys.empty()) {
+ return Status::Invalid("PK real-time writer requires at least one
primary key");
+ }
+ if (restored_max_sequence_number < -1 ||
+ restored_max_sequence_number == std::numeric_limits<int64_t>::max()) {
+ return Status::Invalid("PK restored sequence number is invalid");
+ }
+ arrow::FieldVector key_fields;
+ key_fields.reserve(trimmed_primary_keys.size());
+ for (const std::string& key : trimmed_primary_keys) {
+ std::shared_ptr<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));
+ }
+ arrow::FieldVector prepared_fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())
+ ->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())};
+ prepared_fields.insert(prepared_fields.end(),
write_schema->fields().begin(),
+ write_schema->fields().end());
+ 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,
+ arrow::schema(std::move(prepared_fields)),
arrow::schema(std::move(key_fields)),
+ trimmed_primary_keys, key_comparator, store_state.initial_offset,
+ initial_max_sequence_number, memory_pool));
+}
+
+RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter(
+ const std::shared_ptr<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>& prepared_schema,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator, 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),
+ prepared_schema_(prepared_schema),
+ key_schema_(key_schema),
+ trimmed_primary_keys_(trimmed_primary_keys),
+ key_comparator_(key_comparator),
+ next_offset_(next_offset),
+ last_sequence_number_(last_sequence_number) {}
+
+Status RealtimePrimaryKeyWriter::Write(std::unique_ptr<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> prepared,
+ PrepareBatch(std::move(batch), write_schema_, prepared_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(*prepared,
output.get()));
+ PAIMON_RETURN_NOT_OK(RetainPreparedArrayPool(output.get(), arrow_pool_));
+ RecordBatchBuilder builder(output.get());
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> prepared_batch,
builder.Finish());
+ PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{
+ std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ +
count)}));
+ next_offset_ += count;
+ last_sequence_number_ += count;
+ PAIMON_RETURN_NOT_OK(
+ realtime_context_
+ ->AdvanceMaterializedMaxSequenceNumber(partition_bucket_,
last_sequence_number_)
+ .status());
Review Comment:
`PAIMON_RETURN_NOT_OK(func)`could be enough here;
`PAIMON_RETURN_NOT_OK(func.status())` isn’t necessary.
##########
src/paimon/core/realtime/primary_key_realtime_store.cpp:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 <mutex>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "paimon/common/metrics/metrics_impl.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/core_options.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/schema/table_schema.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const
TableSchema& schema) {
+ if (options.GetBucket() <= 0) {
+ return Status::NotImplemented("PK realtime v1 requires fixed buckets");
+ }
+ if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) {
+ return Status::NotImplemented("PK realtime v1 supports only the
DEDUPLICATE merge engine");
+ }
+ if (options.DataEvolutionEnabled()) {
+ return Status::NotImplemented("PK realtime v1 does not support data
evolution");
+ }
+ if (!options.GetFieldsSequenceGroups().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence groups");
+ }
+ if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete()
||
+ options.AggregationRemoveRecordOnDelete() ||
+ !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) {
+ return Status::NotImplemented("PK realtime v1 requires default delete
behavior");
+ }
+ if (!options.GetSequenceField().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence.field");
+ }
+ if (!options.SequenceFieldSortOrderIsAscending()) {
+ return Status::NotImplemented(
+ "PK realtime v1 supports only ascending
sequence.field.sort-order");
+ }
+ if (options.NeedLookup() || options.DeletionVectorsEnabled() ||
+ options.GetChangelogProducer() != ChangelogProducer::NONE) {
+ return Status::NotImplemented("PK realtime v1 does not support lookup
or early MOR");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::vector<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 v1 does not support FLOAT or DOUBLE primary
keys");
+ }
+ }
+ if (options.GlobalIndexEnabled()) {
+ PAIMON_ASSIGN_OR_RAISE(PrimaryKeyIndexDefinitions definitions,
+ PrimaryKeyIndexDefinitions::Create(schema));
+ if (!definitions.Definitions().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
global indexes");
+ }
+ }
+ return Status::OK();
+}
+
+namespace {
+
+uint64_t GetArrayMemoryUsage(const std::shared_ptr<arrow::ArrayData>& data) {
+ uint64_t total = 0;
+ for (const std::shared_ptr<arrow::Buffer>& buffer : data->buffers) {
+ if (buffer) {
+ total += static_cast<uint64_t>(buffer->size());
+ }
+ }
+ for (const std::shared_ptr<arrow::ArrayData>& child : data->child_data) {
+ total += GetArrayMemoryUsage(child);
+ }
+ if (data->dictionary) {
+ total += GetArrayMemoryUsage(data->dictionary);
+ }
+ return total;
+}
+
+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)
+ : 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>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*data_,
array.get(), schema.get()));
+ data_.reset();
+ return ReadBatch(std::move(array), std::move(schema));
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return metrics_;
+ }
+ void Close() override {
+ data_.reset();
+ }
+
+ private:
+ std::shared_ptr<arrow::StructArray> data_;
+ std::shared_ptr<Metrics> metrics_;
+};
+
+} // namespace
+
+class PrimaryKeyRealtimeStore::Impl {
+ public:
+ explicit Impl(std::shared_ptr<arrow::Schema> prepared_schema)
+ : prepared_schema_(std::move(prepared_schema)) {}
+
+ 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_(prepared_schema_->fields())));
+ if (!array || array->type_id() != arrow::Type::STRUCT) {
+ return Status::Invalid("PK real-time prepared batch is not a
StructArray");
+ }
+ std::shared_ptr<arrow::StructArray> prepared =
+ checked_pointer_cast<arrow::StructArray>(array);
+ std::lock_guard<std::mutex> lock(mutex_);
+ building_.push_back(
+ StoredBatch{prepared, write_batch.offset_range,
GetArrayMemoryUsage(prepared->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));
+ }
+ 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::shared_ptr<RealtimeReadView>(new
ReadView(std::move(segments)));
+ }
+
+ Result<std::vector<std::unique_ptr<BatchReader>>> CreateQueryReaders(
+ const std::shared_ptr<RealtimeReadView>& view, int64_t, const
RealtimeQueryContext&) {
+ 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");
+ }
+ std::vector<std::unique_ptr<BatchReader>> readers;
+ for (const std::shared_ptr<Segment>& segment : typed->Segments()) {
+ for (const StoredBatch& batch : segment->Batches()) {
Review Comment:
My understanding is that we should prune based on the fields in
`RealtimeQueryContext`. Similar to how we set a read schema for the Parquet
format reader, returning all fields and then filtering them at the Paimon layer
doesn’t seem ideal.
##########
src/paimon/core/realtime/prepared_key_value_reader.cpp:
##########
@@ -0,0 +1,698 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "arrow/array/array_base.h"
+#include "arrow/array/array_nested.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/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/reader/batch_reader.h"
+#include "paimon/status.h"
+
+namespace paimon {
+
+namespace {
+
+constexpr int32_t kValueKindIndex = 0;
+constexpr int32_t kSequenceNumberIndex = 1;
+constexpr int32_t kRealtimeOffsetIndex = 2;
+constexpr int32_t kPreparedValueStartIndex = 3;
+
+template <typename Reader>
+void CloseReaders(const std::vector<std::unique_ptr<Reader>>& readers) {
+ for (const std::unique_ptr<Reader>& reader : readers) {
+ if (reader) {
+ reader->Close();
+ }
+ }
+}
+
+Result<std::shared_ptr<arrow::Array>> AlignArrayByPaimonIds(
+ const std::shared_ptr<arrow::Array>& array, const
std::shared_ptr<arrow::DataType>& read_type,
+ arrow::MemoryPool* arrow_pool);
+
+class RealtimeOffsetCoverage {
+ public:
+ static Result<std::shared_ptr<RealtimeOffsetCoverage>> Create(const
OffsetRange& sealed_offsets,
+ size_t
reader_count) {
+ if (sealed_offsets.begin < 0 || sealed_offsets.end <
sealed_offsets.begin) {
+ return Status::Invalid("PK real-time store returned an invalid
sealed offset range");
+ }
+ return std::shared_ptr<RealtimeOffsetCoverage>(
+ new RealtimeOffsetCoverage(sealed_offsets, reader_count));
+ }
+
+ Status Add(const arrow::Int64Array& offsets) {
+ for (int64_t row = 0; row < offsets.length(); ++row) {
+ const int64_t offset = offsets.Value(row);
+ if (offset < sealed_offsets_.begin || offset >=
sealed_offsets_.end) {
+ return Status::Invalid(
+ "PK real-time store commit reader offset is outside the
sealed range");
+ }
+ min_seen_offset_ = std::min(min_seen_offset_, offset);
+ max_seen_offset_ = std::max(max_seen_offset_, offset);
+ ++seen_count_;
+ }
+ return Status::OK();
+ }
Review Comment:
Could we simplify this reader to only apply `_REALTIME_OFFSET` filtering
through `NextBatchWithBitmap()` and then reuse the existing key-value batch
conversion logic?
The store should already return the exact requested schema and handle nested
projection. We could extract the common `BatchReader`-to-`KeyValueRecordReader`
logic from `KeyValueDataFileRecordReader`, make it accept a regular
`BatchReader` with precomputed field indexes, and reuse it for memory readers.
This would remove most of the schema alignment, projection, and
ordering-validation code from `PreparedKeyValueReader`.
I am also fine with treating PK ordering as part of the store plugin
contract without framework-side validation. The current `ValidateOrdering` adds
another full key projection and O(row_count) PK-comparison pass on the hot
path. If runtime validation is required, could it be fused into
`MergedKeyValueRecordReader`, which already compares adjacent keys during
per-reader deduplication, to avoid the duplicated work?
##########
src/paimon/core/mergetree/merge_tree_writer.cpp:
##########
@@ -154,6 +154,59 @@ Status
MergeTreeWriter::Write(std::unique_ptr<RecordBatch>&& moved_batch) {
return Status::OK();
}
+Status MergeTreeWriter::WriteSortedReaders(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& readers) {
Review Comment:
May change func to `FlushSortedReaders` or `WriteSortedReadersToFiles`.
##########
src/paimon/common/table/special_fields.h:
##########
Review Comment:
We may also need to add `RealtimeOffset` here.
##########
src/paimon/core/realtime/prepared_key_value_reader.h:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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 <optional>
+#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 FieldsComparator;
+class MemoryPool;
+
+/// Validates the required leading fields of a prepared real-time transport
schema.
+Status ValidatePreparedTransportSchema(const std::shared_ptr<arrow::Schema>&
prepared_schema);
Review Comment:
Please avoid using global functions in production code when possible. It
would be better to place this in a shared/common/helper class as a static
function instead.
##########
src/paimon/core/realtime/primary_key_realtime_store.h:
##########
@@ -0,0 +1,63 @@
+/*
+ * 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 CoreOptions;
+class TableSchema;
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const
TableSchema& schema);
+
Review Comment:
Similarly, please avoid using global functions here.
##########
src/paimon/core/realtime/realtime_primary_key_writer.cpp:
##########
@@ -0,0 +1,325 @@
+/*
+ * 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/data_field.h"
+#include "paimon/common/types/row_kind.h"
+#include "paimon/common/utils/arrow/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
+#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
+#include "paimon/core/mergetree/merge_tree_writer.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/realtime/realtime_context_impl.h"
+#include "paimon/core/utils/commit_increment.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+namespace {
+
+struct PreparedArrayPrivateData {
+ void (*release)(ArrowArray*);
+ void* private_data;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool;
+};
+
+void ReleasePreparedArray(ArrowArray* array) {
+ auto* data = static_cast<PreparedArrayPrivateData*>(array->private_data);
+ array->release = data->release;
+ array->private_data = data->private_data;
+ array->release(array);
+ delete data;
+}
+
+Status RetainPreparedArrayPool(ArrowArray* array,
+ const std::shared_ptr<arrow::MemoryPool>&
arrow_pool) {
+ if (!array || !array->release || !arrow_pool) {
+ return Status::Invalid("cannot retain prepared batch memory pool");
+ }
+ array->private_data =
+ new PreparedArrayPrivateData{array->release, array->private_data,
arrow_pool};
+ array->release = ReleasePreparedArray;
+ return Status::OK();
+}
+
Review Comment:
This is a very elegant design 👍. We may be able to apply the same idea later
to manage the lifetime of `ArrowArray`s returned by `BatchReader`, so that
`BatchReader` no longer has to outlive the `ArrowArray`.
##########
src/paimon/core/operation/merge_file_split_read.cpp:
##########
@@ -78,6 +78,126 @@ struct KeyValue;
template <typename T>
class MergeFunctionWrapper;
+class MergeFileSplitRead::RealtimeReaderBuilder {
+ public:
+ static Result<std::unique_ptr<BatchReader>> Create(
+ MergeFileSplitRead* owner, const std::vector<std::shared_ptr<Split>>&
disk_splits,
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&&
additional_readers) {
+ 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));
+ }
+ 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::dynamic_pointer_cast<DataSplitImpl>(disk_splits.front());
+ if (!first_split) {
+ return Status::Invalid("merge input disk split is not a data
split");
+ }
+ const BinaryRow& partition = first_split->Partition();
+ const int32_t bucket = first_split->Bucket();
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataFilePathFactory>
data_file_path_factory,
+
owner_->path_factory_->CreateDataFilePathFactory(partition, bucket));
+
+ 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 || !(data_split->Partition() == partition) ||
+ data_split->Bucket() != bucket) {
+ return Status::Invalid("merge input disk splits do not share a
partition-bucket");
+ }
+ if (!data_split->BeforeFiles().empty() ||
data_split->IsStreaming() ||
+ data_split->Bucket() == BucketModeDefine::POSTPONE_BUCKET) {
+ return Status::Invalid("additional merge input requires
fixed-bucket batch splits");
+ }
+ const std::vector<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());
+ }
+ }
+
+ DeletionVector::Factory dv_factory = DeletionVector::CreateFactory(
+ owner_->options_.GetFileSystem(),
+ DeletionVector::CreateDeletionFileMap(data_files, deletion_files),
owner_->pool_);
+ std::vector<std::vector<SortedRun>> disk_sections =
+ IntervalPartition(data_files, owner_->key_comparator_).Partition();
+ for (const std::vector<SortedRun>& section : disk_sections) {
+ for (const SortedRun& run : section) {
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<KeyValueRecordReader>
disk_reader,
+ owner_->CreateReaderForRun(partition,
run, dv_factory,
+
owner_->predicate_for_keys_,
+
data_file_path_factory));
+ readers->push_back(std::move(disk_reader));
+ }
+ }
+ return Status::OK();
+ }
+
+ Result<std::unique_ptr<BatchReader>> CreateMergedReader(
+ std::vector<std::unique_ptr<KeyValueRecordReader>>&& record_readers) {
+ if (record_readers.empty()) {
+ 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)));
+ return CreateProjectedReader(std::move(sort_merge_reader));
+ }
+
+ Result<std::unique_ptr<BatchReader>> CreateProjectedReader(
+ std::unique_ptr<SortMergeReader>&& sort_merge_reader) {
+ if (!owner_->force_keep_delete_) {
+ sort_merge_reader =
std::make_unique<DropDeleteReader>(std::move(sort_merge_reader));
+ }
+
Review Comment:
Could we reduce the duplication here by extracting the common merge-reader
construction steps from `MergeFileSplitRead`?
`CollectDiskReaders` duplicates parts of the existing disk merge path,
including deletion-vector factory creation, interval partitioning, and
per-`SortedRun` reader construction. Meanwhile,
`RealtimeReaderBuilder::CreateProjectedReader` duplicates the existing
drop-delete, sync/async projection, predicate-filtering, and row-kind
completion pipeline.
For example, we could extract reusable helpers such as
`CreateRecordReadersForSection(...)` and `CreateProjectedReader(...)`, and use
them from both the normal and real-time paths. The real-time builder would then
only be responsible for validating and flattening multiple disk splits and
combining the resulting disk readers with memory readers.
This would make the real-time implementation much smaller and reduce the
risk of the normal and real-time MOR pipelines diverging when projection,
predicate handling, or reader construction changes.
##########
src/paimon/core/realtime/primary_key_realtime_store.cpp:
##########
@@ -0,0 +1,327 @@
+/*
+ * 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 <mutex>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "paimon/common/metrics/metrics_impl.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/core_options.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/schema/table_schema.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const
TableSchema& schema) {
+ if (options.GetBucket() <= 0) {
+ return Status::NotImplemented("PK realtime v1 requires fixed buckets");
+ }
+ if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) {
+ return Status::NotImplemented("PK realtime v1 supports only the
DEDUPLICATE merge engine");
+ }
+ if (options.DataEvolutionEnabled()) {
+ return Status::NotImplemented("PK realtime v1 does not support data
evolution");
+ }
+ if (!options.GetFieldsSequenceGroups().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence groups");
+ }
+ if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete()
||
+ options.AggregationRemoveRecordOnDelete() ||
+ !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) {
+ return Status::NotImplemented("PK realtime v1 requires default delete
behavior");
+ }
+ if (!options.GetSequenceField().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence.field");
+ }
+ if (!options.SequenceFieldSortOrderIsAscending()) {
+ return Status::NotImplemented(
+ "PK realtime v1 supports only ascending
sequence.field.sort-order");
+ }
+ if (options.NeedLookup() || options.DeletionVectorsEnabled() ||
+ options.GetChangelogProducer() != ChangelogProducer::NONE) {
+ return Status::NotImplemented("PK realtime v1 does not support lookup
or early MOR");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::vector<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 v1 does not support FLOAT or DOUBLE primary
keys");
+ }
Review Comment:
I’m a bit curious — why don’t we support float and double as primary keys?
##########
src/paimon/core/realtime/realtime_context_impl.cpp:
##########
@@ -117,27 +140,48 @@ 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_ASSIGN_OR_RAISE(
- std::shared_ptr<RealtimeStore> store,
- factory_->Create(std::move(write_schema), statistics_mode, options,
memory_pool));
- stores_.emplace(key, store);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+ arrow::ExportSchema(*requested_schema, request.write_schema.get()));
+ RealtimeStoreMode mode = request.mode;
+ Result<std::shared_ptr<RealtimeStore>> store_result =
factory_->Create(std::move(request));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeStore> store,
std::move(store_result));
Review Comment:
Why was this changed from directly using `PAIMON_ASSIGN_OR_RAISE(...)` to
first storing the result in `Result<std::shared_ptr<RealtimeStore>>
store_result` and then unwrapping it?
##########
src/paimon/core/realtime/realtime_primary_key_writer.cpp:
##########
@@ -0,0 +1,325 @@
+/*
+ * 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/data_field.h"
+#include "paimon/common/types/row_kind.h"
+#include "paimon/common/utils/arrow/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/mergetree/compact/deduplicate_merge_function.h"
+#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h"
+#include "paimon/core/mergetree/merge_tree_writer.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/realtime/realtime_context_impl.h"
+#include "paimon/core/utils/commit_increment.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+namespace {
+
+struct PreparedArrayPrivateData {
+ void (*release)(ArrowArray*);
+ void* private_data;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool;
+};
+
+void ReleasePreparedArray(ArrowArray* array) {
+ auto* data = static_cast<PreparedArrayPrivateData*>(array->private_data);
+ array->release = data->release;
+ array->private_data = data->private_data;
+ array->release(array);
+ delete data;
+}
+
+Status RetainPreparedArrayPool(ArrowArray* array,
+ const std::shared_ptr<arrow::MemoryPool>&
arrow_pool) {
+ if (!array || !array->release || !arrow_pool) {
+ return Status::Invalid("cannot retain prepared batch memory pool");
+ }
+ array->private_data =
+ new PreparedArrayPrivateData{array->release, array->private_data,
arrow_pool};
+ array->release = ReleasePreparedArray;
+ return Status::OK();
+}
+
+Result<std::shared_ptr<arrow::StructArray>> PrepareBatch(
+ std::unique_ptr<RecordBatch>&& batch, const
std::shared_ptr<arrow::Schema>& write_schema,
+ const std::shared_ptr<arrow::Schema>& prepared_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> prepared,
+ arrow::StructArray::Make(std::move(columns),
prepared_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(prepared), options,
&context));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum sorted,
+ arrow::compute::Take(arrow::Datum(prepared), indices,
+ arrow::compute::TakeOptions::NoBoundsCheck(),
&context));
+ std::shared_ptr<arrow::Array> sorted_array = sorted.make_array();
+ if (!sorted_array || sorted_array->type_id() != arrow::Type::STRUCT) {
+ return Status::Invalid("PK real-time sorted batch is not a
StructArray");
+ }
+ return checked_pointer_cast<arrow::StructArray>(std::move(sorted_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::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ 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 (!store_state.store || !merge_tree_writer || !write_schema ||
!key_comparator ||
+ !realtime_context || !memory_pool) {
+ return Status::Invalid("PK real-time writer received a null
dependency");
+ }
+ if (trimmed_primary_keys.empty()) {
+ return Status::Invalid("PK real-time writer requires at least one
primary key");
+ }
+ if (restored_max_sequence_number < -1 ||
+ restored_max_sequence_number == std::numeric_limits<int64_t>::max()) {
+ return Status::Invalid("PK restored sequence number is invalid");
+ }
+ arrow::FieldVector key_fields;
+ key_fields.reserve(trimmed_primary_keys.size());
+ for (const std::string& key : trimmed_primary_keys) {
+ std::shared_ptr<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));
+ }
+ arrow::FieldVector prepared_fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber())
+ ->WithNullable(false),
+
DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset())};
+ prepared_fields.insert(prepared_fields.end(),
write_schema->fields().begin(),
+ write_schema->fields().end());
+ 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,
+ arrow::schema(std::move(prepared_fields)),
arrow::schema(std::move(key_fields)),
+ trimmed_primary_keys, key_comparator, store_state.initial_offset,
+ initial_max_sequence_number, memory_pool));
+}
+
+RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter(
+ const std::shared_ptr<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>& prepared_schema,
+ const std::shared_ptr<arrow::Schema>& key_schema,
+ const std::vector<std::string>& trimmed_primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator, 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),
+ prepared_schema_(prepared_schema),
+ key_schema_(key_schema),
+ trimmed_primary_keys_(trimmed_primary_keys),
+ key_comparator_(key_comparator),
+ next_offset_(next_offset),
+ last_sequence_number_(last_sequence_number) {}
+
+Status RealtimePrimaryKeyWriter::Write(std::unique_ptr<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> prepared,
+ PrepareBatch(std::move(batch), write_schema_, prepared_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(*prepared,
output.get()));
+ PAIMON_RETURN_NOT_OK(RetainPreparedArrayPool(output.get(), arrow_pool_));
+ RecordBatchBuilder builder(output.get());
+ PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> prepared_batch,
builder.Finish());
+ PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{
+ std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ +
count)}));
+ next_offset_ += count;
+ last_sequence_number_ += count;
+ PAIMON_RETURN_NOT_OK(
+ realtime_context_
+ ->AdvanceMaterializedMaxSequenceNumber(partition_bucket_,
last_sequence_number_)
+ .status());
Review Comment:
`materialized_max_sequence_number` seems to be used only to support writer
handoff with uncommitted state retained in the same `RealtimeContext`. For
failure recovery, the old context should be discarded, and the sequence number
should be restored from the latest snapshot before replaying the data. Reusing
this value would instead assign larger sequence numbers during replay. Could
you confirm whether handoff with uncommitted state is a required use case? If
not, I suggest removing this state and always restoring the sequence number
from files.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
