HaHaJeff commented on code in PR #224:
URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3852187849
##########
src/paimon/core/realtime/realtime_context_impl.cpp:
##########
@@ -78,27 +78,37 @@ Status RealtimeContextImpl::Start() {
}
Result<RealtimeStoreState> RealtimeContextImpl::GetOrCreateRealtimeStore(
- 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) {
+ RealtimeStoreCreateRequest&& request) {
std::lock_guard<std::mutex> progress_lock(progress_mutex_);
std::lock_guard<std::mutex> registry_lock(mutex_);
- const RealtimePartitionBucket key(partition, bucket);
+ const RealtimePartitionBucket key(request.partition, request.bucket);
+ std::optional<int64_t> initial_max_sequence_number;
+ PrimaryKeyRealtimeStoreCreateConfig* primary_key_config =
+ std::get_if<PrimaryKeyRealtimeStoreCreateConfig>(&request.mode_config);
+ if (primary_key_config) {
+ auto [sequence_iter, inserted] =
materialized_max_sequence_numbers_.emplace(
+ key, primary_key_config->restore_max_sequence_number);
+ if (!inserted && primary_key_config->restore_max_sequence_number >
sequence_iter->second) {
Review Comment:
Thanks for catching this. The original implementation let the store assign
sequence numbers because it also performed the PK merge. After moving
preparation and MOR into the framework, keeping another allocator in the store
would create two sequence authorities and make writer handoff unsafe. Commit
141099c7e4f57e3cdf40d3c865619f61a8b8e7f8 moved the partition-bucket
materialized sequence watermark into RealtimeContext; replacement writers
initialize from it and advance it only after a successful store write. V1
follows the same lifecycle contract as append realtime: one RealtimeContext is
owned by one active FileStoreWrite, so simultaneous writers sharing a context
are unsupported.
##########
src/paimon/core/realtime/primary_key_realtime_store.cpp:
##########
@@ -0,0 +1,597 @@
+/*
+ * 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 <limits>
+#include <mutex>
+#include <utility>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "paimon/common/data/binary_row_writer.h"
+#include "paimon/common/data/columnar/columnar_row_ref.h"
+#include "paimon/common/metrics/metrics_impl.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/common/utils/fields_comparator.h"
+#include "paimon/core/core_options.h"
+#include "paimon/core/io/key_value_in_memory_record_reader.h"
+#include "paimon/core/io/key_value_projection_consumer.h"
+#include "paimon/core/io/key_value_projection_reader.h"
+#include "paimon/core/io/merged_key_value_record_reader.h"
+#include "paimon/core/key_value.h"
+#include "paimon/core/mergetree/compact/sort_merge_reader_with_loser_tree.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options) {
+ 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");
+ }
+ return Status::OK();
+}
+
+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;
+}
+
+struct StoredBatch {
+ std::shared_ptr<arrow::StructArray> data;
+ std::vector<RecordBatch::RowKind> row_kinds;
+ OffsetRange offset_range;
+ int64_t first_sequence_number;
+ uint64_t memory_usage;
+};
+using BatchGroup = std::vector<std::shared_ptr<const StoredBatch>>;
+
+class Segment final : public RealtimeSegmentHandle {
+ public:
+ Segment(const OffsetRange& offset_range,
+ std::vector<std::shared_ptr<const StoredBatch>>&& batches)
+ : offset_range_(offset_range), batches_(std::move(batches)) {}
+
+ OffsetRange GetOffsetRange() const override {
+ return offset_range_;
+ }
+
+ const std::vector<std::shared_ptr<const StoredBatch>>& Batches() const {
+ return batches_;
+ }
+
+ uint64_t GetMemoryUsage() const {
+ uint64_t result = 0;
+ for (const std::shared_ptr<const StoredBatch>& batch : batches_) {
+ result += batch->memory_usage;
+ }
+ return result;
+ }
+
+ private:
+ OffsetRange offset_range_;
+ std::vector<std::shared_ptr<const StoredBatch>> batches_;
+};
+
+class PrimaryKeyRealtimeReadView final : public RealtimeReadView {
+ public:
+ explicit PrimaryKeyRealtimeReadView(std::vector<BatchGroup>&& groups)
+ : groups_(std::move(groups)) {
+ if (!groups_.empty()) {
+ offset_range_ =
OffsetRange(groups_.front().front()->offset_range.begin,
+
groups_.back().back()->offset_range.end);
+ }
+ }
+
+ std::optional<OffsetRange> GetOffsetRange() const override {
+ return offset_range_;
+ }
+
+ const std::vector<BatchGroup>& Groups() const {
+ return groups_;
+ }
+
+ private:
+ std::vector<BatchGroup> groups_;
+ std::optional<OffsetRange> offset_range_;
+};
+
+class CommitBatchReader final : public BatchReader {
+ public:
+ CommitBatchReader(const std::shared_ptr<Segment>& segment,
+ const std::shared_ptr<arrow::MemoryPool>& arrow_pool)
+ : segment_(segment), arrow_pool_(arrow_pool),
metrics_(std::make_shared<MetricsImpl>()) {}
+
+ Result<ReadBatch> NextBatch() override {
+ if (!segment_ || next_batch_ >=
static_cast<int32_t>(segment_->Batches().size())) {
+ return MakeEofBatch();
+ }
+ const std::shared_ptr<const StoredBatch>& stored =
segment_->Batches()[next_batch_++];
+ const int64_t row_count = stored->data->length();
+ arrow::Int8Builder row_kind_builder(arrow_pool_.get());
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Reserve(row_count));
+ if (stored->row_kinds.empty()) {
+ for (int64_t i = 0; i < row_count; ++i) {
+
row_kind_builder.UnsafeAppend(static_cast<int8_t>(RecordBatch::RowKind::INSERT));
+ }
+ } else {
+ for (RecordBatch::RowKind row_kind : stored->row_kinds) {
+ row_kind_builder.UnsafeAppend(static_cast<int8_t>(row_kind));
+ }
+ }
+ std::shared_ptr<arrow::Array> row_kind_array;
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(row_kind_builder.Finish(&row_kind_array));
+ arrow::ArrayVector arrays = {std::move(row_kind_array)};
+ arrays.insert(arrays.end(), stored->data->fields().begin(),
stored->data->fields().end());
+ arrow::FieldVector fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())};
+ const arrow::FieldVector& value_fields =
stored->data->struct_type()->fields();
+ fields.insert(fields.end(), value_fields.begin(), value_fields.end());
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray>
output,
+ arrow::StructArray::Make(arrays,
fields));
+ auto c_array = std::make_unique<ArrowArray>();
+ auto c_schema = std::make_unique<ArrowSchema>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*output,
c_array.get(), c_schema.get()));
+ return ReadBatch(std::move(c_array), std::move(c_schema));
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return metrics_;
+ }
+
+ void Close() override {
+ segment_.reset();
+ }
+
+ private:
+ std::shared_ptr<Segment> segment_;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+ std::shared_ptr<Metrics> metrics_;
+ int32_t next_batch_ = 0;
+};
+
+class KeyRangeBatchReader final : public BatchReader, public
PrimaryKeyRangeProvider {
+ public:
+ KeyRangeBatchReader(std::unique_ptr<BatchReader>&& reader,
+ const std::shared_ptr<InternalRow>& min_key,
+ const std::shared_ptr<InternalRow>& max_key)
+ : reader_(std::move(reader)), min_key_(min_key), max_key_(max_key) {}
+
+ Result<ReadBatch> NextBatch() override {
+ return reader_->NextBatch();
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return reader_->GetReaderMetrics();
+ }
+
+ void Close() override {
+ reader_->Close();
+ }
+
+ std::shared_ptr<InternalRow> GetMinKey() const override {
+ return min_key_;
+ }
+
+ std::shared_ptr<InternalRow> GetMaxKey() const override {
+ return max_key_;
+ }
+
+ private:
+ std::unique_ptr<BatchReader> reader_;
+ std::shared_ptr<InternalRow> min_key_;
+ std::shared_ptr<InternalRow> max_key_;
+};
+
+} // namespace
+
+class PrimaryKeyRealtimeStore::Impl {
+ public:
+ Impl(const std::shared_ptr<arrow::Schema>& write_schema,
std::vector<std::string> primary_keys,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ const
std::function<std::shared_ptr<MergeFunctionWrapper<KeyValue>>()>&
+ merge_function_wrapper_factory,
+ int64_t next_sequence_number, int32_t read_batch_size,
+ const std::shared_ptr<MemoryPool>& memory_pool)
+ : write_schema_(write_schema),
+ primary_keys_(std::move(primary_keys)),
+ key_comparator_(key_comparator),
+ merge_function_wrapper_factory_(merge_function_wrapper_factory),
+ next_sequence_number_(next_sequence_number),
+ read_batch_size_(read_batch_size),
+ memory_pool_(memory_pool),
+ arrow_pool_(GetArrowPool(memory_pool)) {}
+
+ Result<std::shared_ptr<InternalRow>> CopyKey(const InternalRow& key) const
{
+ auto result =
std::make_shared<BinaryRow>(static_cast<int32_t>(primary_keys_.size()));
+ BinaryRowWriter writer(result.get(), /*initial_size=*/128,
memory_pool_.get());
+ writer.Reset();
+ for (int32_t index = 0; index <
static_cast<int32_t>(primary_keys_.size()); ++index) {
+ std::shared_ptr<arrow::Field> field =
+ write_schema_->GetFieldByName(primary_keys_[index]);
+ PAIMON_ASSIGN_OR_RAISE(InternalRow::FieldGetterFunc getter,
+ InternalRow::CreateFieldGetter(index,
field->type(),
+
/*use_view=*/true));
+ PAIMON_ASSIGN_OR_RAISE(BinaryRowWriter::FieldSetterFunc setter,
+ BinaryRowWriter::CreateFieldSetter(index,
field->type()));
+ setter(getter(key), &writer);
+ }
+ writer.Complete();
+ return std::static_pointer_cast<InternalRow>(result);
+ }
+
+ Result<std::pair<std::shared_ptr<InternalRow>,
std::shared_ptr<InternalRow>>> GetKeyRange(
+ const std::shared_ptr<arrow::StructArray>& values) const {
+ arrow::ArrayVector key_arrays;
+ key_arrays.reserve(primary_keys_.size());
+ for (const std::string& primary_key : primary_keys_) {
+ std::shared_ptr<arrow::Array> key_array =
values->GetFieldByName(primary_key);
+ if (!key_array) {
+ return Status::Invalid("primary key is missing from PK query
batch: ", primary_key);
+ }
+ key_arrays.push_back(std::move(key_array));
+ }
+ auto context = std::make_shared<ColumnarBatchContext>(key_arrays,
memory_pool_);
+ int64_t min_row = 0;
+ int64_t max_row = 0;
+ for (int64_t row = 1; row < values->length(); ++row) {
+ ColumnarRowRef current(context, row);
+ ColumnarRowRef min_key(context, min_row);
+ ColumnarRowRef max_key(context, max_row);
+ if (key_comparator_->CompareTo(current, min_key) < 0) {
+ min_row = row;
+ }
+ if (key_comparator_->CompareTo(current, max_key) > 0) {
+ max_row = row;
+ }
+ }
+ ColumnarRowRef min_key(context, min_row);
+ ColumnarRowRef max_key(context, max_row);
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InternalRow> copied_min,
CopyKey(min_key));
+ PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InternalRow> copied_max,
CopyKey(max_key));
+ return std::make_pair(std::move(copied_min), std::move(copied_max));
+ }
+
+ 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 (row_count <= 0 || write_batch.offset_range.begin < 0 ||
+ write_batch.offset_range.Count() != row_count) {
+ return Status::Invalid("PK real-time offset range does not match
batch row count");
+ }
+ const std::vector<RecordBatch::RowKind>& row_kinds =
write_batch.batch->GetRowKind();
+ if (!row_kinds.empty() && static_cast<int64_t>(row_kinds.size()) !=
row_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);
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Array> imported,
+ arrow::ImportArray(write_batch.batch->GetData(),
+ arrow::struct_(write_schema_->fields())));
+ if (!imported || imported->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>(imported);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(values->ValidateFull());
+
+ std::lock_guard<std::mutex> lock(mutex_);
+ if (last_offset_ && write_batch.offset_range.begin !=
last_offset_.value()) {
+ return Status::Invalid("PK real-time offset ranges must be
contiguous");
+ }
+ if (row_count > std::numeric_limits<int64_t>::max() -
next_sequence_number_) {
+ return Status::Invalid("PK sequence range exceeds INT64_MAX");
+ }
+ auto stored = std::make_shared<const StoredBatch>(
+ StoredBatch{std::move(values), row_kinds, write_batch.offset_range,
+ next_sequence_number_,
GetArrayMemoryUsage(imported->data())});
+ building_batches_.push_back(std::move(stored));
+ building_memory_usage_ += building_batches_.back()->memory_usage;
+ last_offset_ = write_batch.offset_range.end;
+ next_sequence_number_ += row_count;
+ return Status::OK();
+ }
+
+ Result<std::optional<std::shared_ptr<RealtimeSegmentHandle>>>
SealForCommit() {
+ std::lock_guard<std::mutex> lock(mutex_);
+ if (building_batches_.empty()) {
+ return std::optional<std::shared_ptr<RealtimeSegmentHandle>>();
+ }
+ const OffsetRange range(building_batches_.front()->offset_range.begin,
+ building_batches_.back()->offset_range.end);
+ auto segment = std::make_shared<Segment>(range,
std::move(building_batches_));
+ sealed_segments_.push_back(segment);
+ building_batches_.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>& segment) {
+ std::shared_ptr<Segment> typed =
std::dynamic_pointer_cast<Segment>(segment);
+ if (!typed) {
+ return Status::Invalid("segment was not created by the PK
real-time store");
+ }
+ std::vector<std::unique_ptr<BatchReader>> result;
+ result.push_back(std::make_unique<CommitBatchReader>(typed,
arrow_pool_));
+ return result;
+ }
+
+ Result<std::shared_ptr<RealtimeReadView>> AcquireReadView() {
+ std::lock_guard<std::mutex> lock(mutex_);
+ std::vector<BatchGroup> groups;
+ groups.reserve(sealed_segments_.size() + (building_batches_.empty() ?
0 : 1));
+ for (const std::shared_ptr<Segment>& segment : sealed_segments_) {
+ groups.push_back(segment->Batches());
+ }
+ if (!building_batches_.empty()) {
+ groups.push_back(building_batches_);
+ }
+ return std::shared_ptr<RealtimeReadView>(new
PrimaryKeyRealtimeReadView(std::move(groups)));
+ }
+
+ Result<std::vector<std::unique_ptr<BatchReader>>> CreateQueryReaders(
+ const std::shared_ptr<RealtimeReadView>& view, int64_t lower,
+ const RealtimeQueryContext& context) {
+ std::shared_ptr<PrimaryKeyRealtimeReadView> typed =
+ std::dynamic_pointer_cast<PrimaryKeyRealtimeReadView>(view);
+ if (!typed) {
+ return Status::Invalid("read view was not created by the PK
real-time store");
+ }
+ if (!context.read_schema || !context.read_schema->release) {
+ return Status::Invalid("PK real-time query read schema is null");
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema>
requested,
+
arrow::ImportSchema(context.read_schema));
+ arrow::FieldVector output_fields = {
+
DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())};
+ std::vector<int32_t> projection =
{KeyValueProjectionConsumer::kValueKindProjection};
+ for (const std::shared_ptr<arrow::Field>& field : requested->fields())
{
+ if (field->name() == SpecialFields::ValueKind().Name()) {
+ continue;
+ }
+ output_fields.push_back(field);
+ if (field->name() == SpecialFields::SequenceNumber().Name()) {
+
projection.push_back(KeyValueProjectionConsumer::kSequenceNumberProjection);
+ continue;
+ }
+ const int32_t index = write_schema_->GetFieldIndex(field->name());
Review Comment:
Addressed in df322c1d38dc7e30bd2ca44c4d85297eddf96aac, with missing-field
and plugin-schema handling hardened in
8ab981752c64242cee80cfab66d99c17dbd2febd. The framework adapter recursively
aligns structs, lists, and maps by Paimon field ID before building key/value
readers, so nested child order no longer depends on stored struct position.
Unit and disk-plus-memory integration tests cover this path.
##########
include/paimon/realtime/realtime_store.h:
##########
@@ -40,6 +42,29 @@ namespace paimon {
class MemoryPool;
class Predicate;
+struct PAIMON_EXPORT AppendRealtimeStoreCreateConfig {};
+
+struct PAIMON_EXPORT PrimaryKeyRealtimeStoreCreateConfig {
+ std::vector<std::string> primary_keys;
+ /// Largest sequence restored from the committed snapshot. A PK store
assigns one contiguous
+ /// sequence to every mutation in `Write` order, starting at the next
value, and rejects
+ /// `Write` before the assigned sequence would exceed `INT64_MAX - 1`.
+ int64_t restore_max_sequence_number;
+};
+
+using RealtimeStoreCreateConfig =
+ std::variant<AppendRealtimeStoreCreateConfig,
PrimaryKeyRealtimeStoreCreateConfig>;
+
+struct PAIMON_EXPORT RealtimeStoreCreateRequest {
+ /// Complete table write schema whose ownership is transferred to the
factory.
+ std::unique_ptr<::ArrowSchema> write_schema;
+ std::map<std::string, std::string> options;
+ std::shared_ptr<MemoryPool> memory_pool;
+ std::map<std::string, std::string> partition;
+ int32_t bucket = -1;
+ RealtimeStoreCreateConfig mode_config;
+};
+
Review Comment:
Thanks for the design guidance. The initial implementation placed sequence
assignment, sorting, and in-memory MOR in the PK store, which made custom
stores understand Paimon merge semantics and caused data to be sorted again
during file writing. Commits 8d961530efa9ce41b8885d9ed17160ac5dcdcaa2,
df322c1d38dc7e30bd2ca44c4d85297eddf96aac, and
87e454803670fa909fb5a711bca657d2375918fa moved sorted-reader writing, reader
adaptation, transport-field materialization, and PK sorting into the framework.
Commit ff444191efa75ec4fe292dfb52081f01121135d9 finalized the PK storage-only
boundary by removing store-side PK merge dependencies. The framework allocates
sequence and offsets under the V1 one-context/one-active-writer contract. PK
statistics pruning remains follow-up work; predicate pruning before MOR stays
disabled for correctness.
##########
src/paimon/core/realtime/primary_key_realtime_store.cpp:
##########
@@ -0,0 +1,509 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "paimon/core/realtime/primary_key_realtime_store.h"
+
+#include <cstddef>
+#include <mutex>
+#include <optional>
+#include <queue>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "arrow/api.h"
+#include "arrow/c/bridge.h"
+#include "arrow/compute/api.h"
+#include "paimon/common/data/columnar/columnar_batch_context.h"
+#include "paimon/common/data/columnar/columnar_row_ref.h"
+#include "paimon/common/metrics/metrics_impl.h"
+#include "paimon/common/types/data_field.h"
+#include "paimon/common/utils/arrow/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/core/core_options.h"
+#include "paimon/core/index/pk/primary_key_index_definitions.h"
+#include "paimon/core/realtime/prepared_key_value_reader.h"
+#include "paimon/core/schema/table_schema.h"
+#include "paimon/macros.h"
+
+namespace paimon {
+
+Status ValidatePrimaryKeyRealtimeOptions(const CoreOptions& options, const
TableSchema& schema) {
+ if (options.GetBucket() <= 0) {
+ return Status::NotImplemented("PK realtime v1 requires fixed buckets");
+ }
+ if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) {
+ return Status::NotImplemented("PK realtime v1 supports only the
DEDUPLICATE merge engine");
+ }
+ if (options.DataEvolutionEnabled()) {
+ return Status::NotImplemented("PK realtime v1 does not support data
evolution");
+ }
+ if (!options.GetFieldsSequenceGroups().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence groups");
+ }
+ if (options.IgnoreDelete() || options.PartialUpdateRemoveRecordOnDelete()
||
+ options.AggregationRemoveRecordOnDelete() ||
+ !options.GetPartialUpdateRemoveRecordOnSequenceGroup().empty()) {
+ return Status::NotImplemented("PK realtime v1 requires default delete
behavior");
+ }
+ if (!options.GetSequenceField().empty()) {
+ return Status::NotImplemented("PK realtime v1 does not support
sequence.field");
+ }
+ if (!options.SequenceFieldSortOrderIsAscending()) {
+ return Status::NotImplemented(
+ "PK realtime v1 supports only ascending
sequence.field.sort-order");
+ }
+ if (options.NeedLookup() || options.DeletionVectorsEnabled() ||
+ options.GetChangelogProducer() != ChangelogProducer::NONE) {
+ return Status::NotImplemented("PK realtime v1 does not support lookup
or early MOR");
+ }
+ PAIMON_ASSIGN_OR_RAISE(std::vector<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 RawBatchReader final : public BatchReader {
+ public:
+ RawBatchReader(std::vector<StoredBatch> batches, std::vector<int32_t>
key_field_indexes,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ const std::shared_ptr<MemoryPool>& memory_pool)
+ : batches_(std::move(batches)),
+ positions_(batches_.size(), 0),
+ key_field_indexes_(std::move(key_field_indexes)),
+ key_comparator_(key_comparator),
+ memory_pool_(memory_pool),
+ arrow_pool_(GetArrowPool(memory_pool)),
+ heap_(SourceGreater{this}),
+ metrics_(std::make_shared<MetricsImpl>()) {
+ key_contexts_.reserve(batches_.size());
+ sequence_arrays_.reserve(batches_.size());
+ for (size_t i = 0; i < batches_.size(); ++i) {
+ const StoredBatch& batch = batches_[i];
+ arrow::ArrayVector key_arrays;
+ key_arrays.reserve(key_field_indexes_.size());
+ for (int32_t field_index : key_field_indexes_) {
+ key_arrays.push_back(batch.data->field(field_index));
+ }
+ key_contexts_.push_back(
+ std::make_shared<ColumnarBatchContext>(key_arrays,
memory_pool_));
+ sequence_arrays_.push_back(
+ checked_pointer_cast<arrow::Int64Array>(batch.data->field(1)));
+ if (batch.data->length() > 0) {
+ heap_.push(i);
+ }
+ }
+ }
+
+ Result<ReadBatch> NextBatch() override {
+ if (heap_.empty()) {
+ return MakeEofBatch();
+ }
+
+ struct SelectedRow {
+ size_t selected_source;
+ int64_t source_ordinal;
+ };
+ struct SelectedSource {
+ size_t source;
+ std::vector<int64_t> rows;
+ int64_t base = -1;
+ };
+ std::vector<SelectedRow> selected_rows;
+ selected_rows.reserve(kOutputBatchSize);
+ std::vector<SelectedSource> selected_sources;
+ std::unordered_map<size_t, size_t> selected_source_indexes;
+ while (!heap_.empty() && selected_rows.size() < kOutputBatchSize) {
+ const size_t source = heap_.top();
+ heap_.pop();
+ auto [source_it, inserted] =
+ selected_source_indexes.emplace(source,
selected_sources.size());
+ if (inserted) {
+ selected_sources.push_back(SelectedSource{source, {}});
+ }
+ SelectedSource& selected_source =
selected_sources[source_it->second];
+ selected_rows.push_back(
+ SelectedRow{source_it->second,
static_cast<int64_t>(selected_source.rows.size())});
+ selected_source.rows.push_back(positions_[source]++);
+ if (positions_[source] < batches_[source].data->length()) {
+ heap_.push(source);
+ }
+ }
+
+ arrow::compute::ExecContext context(arrow_pool_.get());
+ arrow::ArrayVector grouped_batches;
+ int64_t grouped_row_count = 0;
+ for (SelectedSource& selected_source : selected_sources) {
+ arrow::Int64Builder source_index_builder(arrow_pool_.get());
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(
+ source_index_builder.AppendValues(selected_source.rows));
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array>
source_indices,
+ source_index_builder.Finish());
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum source_batch,
+
arrow::compute::Take(arrow::Datum(batches_[selected_source.source].data),
+ arrow::Datum(source_indices),
+
arrow::compute::TakeOptions::NoBoundsCheck(), &context));
+ selected_source.base = grouped_row_count;
+ grouped_row_count +=
static_cast<int64_t>(selected_source.rows.size());
+ grouped_batches.push_back(source_batch.make_array());
+ }
+
+ std::shared_ptr<arrow::Array> batch;
+ if (grouped_batches.size() == 1) {
+ batch = std::move(grouped_batches[0]);
+ } else {
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Array> grouped,
+ arrow::Concatenate(grouped_batches, arrow_pool_.get()));
+ arrow::Int64Builder order_builder(arrow_pool_.get());
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(order_builder.Reserve(selected_rows.size()));
+ for (const SelectedRow& selected : selected_rows) {
+
order_builder.UnsafeAppend(selected_sources[selected.selected_source].base +
+ selected.source_ordinal);
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array>
order,
+ order_builder.Finish());
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum reordered,
+ arrow::compute::Take(arrow::Datum(grouped),
arrow::Datum(order),
+
arrow::compute::TakeOptions::NoBoundsCheck(), &context));
+ batch = reordered.make_array();
+ }
+ auto array = std::make_unique<ArrowArray>();
+ auto schema = std::make_unique<ArrowSchema>();
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch,
array.get(), schema.get()));
+ return ReadBatch(std::move(array), std::move(schema));
+ }
+
+ std::shared_ptr<Metrics> GetReaderMetrics() const override {
+ return metrics_;
+ }
+ void Close() override {
+ while (!heap_.empty()) {
+ heap_.pop();
+ }
+ batches_.clear();
+ positions_.clear();
+ key_contexts_.clear();
+ sequence_arrays_.clear();
+ }
+
+ private:
+ static constexpr size_t kOutputBatchSize = 1024;
+
+ bool Less(size_t left, size_t right) const {
+ ColumnarRowRef left_key(key_contexts_[left], positions_[left]);
+ ColumnarRowRef right_key(key_contexts_[right], positions_[right]);
+ const int32_t key_comparison = key_comparator_->CompareTo(left_key,
right_key);
+ if (key_comparison != 0) {
+ return key_comparison < 0;
+ }
+ const int64_t left_sequence =
sequence_arrays_[left]->Value(positions_[left]);
+ const int64_t right_sequence =
sequence_arrays_[right]->Value(positions_[right]);
+ if (left_sequence != right_sequence) {
+ return left_sequence < right_sequence;
+ }
+ return left < right;
+ }
+
+ struct SourceGreater {
+ RawBatchReader* reader;
+
+ bool operator()(size_t left, size_t right) const {
+ return reader->Less(right, left);
+ }
+ };
+
+ std::vector<StoredBatch> batches_;
+ std::vector<int64_t> positions_;
+ std::vector<int32_t> key_field_indexes_;
+ std::shared_ptr<FieldsComparator> key_comparator_;
+ std::shared_ptr<MemoryPool> memory_pool_;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+ std::vector<std::shared_ptr<ColumnarBatchContext>> key_contexts_;
+ std::vector<std::shared_ptr<arrow::Int64Array>> sequence_arrays_;
+ std::priority_queue<size_t, std::vector<size_t>, SourceGreater> heap_;
+ std::shared_ptr<Metrics> metrics_;
+};
+
+} // namespace
+
+class PrimaryKeyRealtimeStore::Impl {
+ public:
+ Impl(std::shared_ptr<arrow::Schema> prepared_schema, std::vector<int32_t>
key_field_indexes,
+ const std::shared_ptr<FieldsComparator>& key_comparator,
+ const std::shared_ptr<MemoryPool>& memory_pool)
+ : prepared_schema_(std::move(prepared_schema)),
+ key_field_indexes_(std::move(key_field_indexes)),
+ key_comparator_(key_comparator),
+ memory_pool_(memory_pool) {}
+
+ Status Write(RealtimeWriteBatch&& write_batch) {
+ if (!write_batch.batch || !write_batch.batch->GetData()) {
+ return Status::Invalid("PK real-time write batch is null");
+ }
+ const int64_t row_count = write_batch.batch->GetData()->length;
+ if (write_batch.offset_range.begin < 0 ||
write_batch.offset_range.Count() != row_count ||
+ row_count <= 0) {
+ return Status::Invalid("PK real-time offset range does not match
batch row count");
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<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);
+ PAIMON_RETURN_NOT_OK_FROM_ARROW(prepared->ValidateFull());
Review Comment:
Thanks, agreed. The redundant ValidateFull() was removed in
ff444191efa75ec4fe292dfb52081f01121135d9. The built-in store now imports and
retains the framework-prepared batch without another full-array validation pass
after the framework kernels have consumed it.
##########
src/paimon/core/realtime/prepared_key_value_reader.cpp:
##########
@@ -0,0 +1,771 @@
+/*
+ * 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 <cstdint>
+#include <memory>
+#include <mutex>
+#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/array/builder_primitive.h"
+#include "arrow/buffer.h"
+#include "arrow/c/bridge.h"
+#include "arrow/compute/api.h"
+#include "arrow/type.h"
+#include "arrow/util/bit_util.h"
+#include "fmt/format.h"
+#include "paimon/common/data/columnar/columnar_batch_context.h"
+#include "paimon/common/data/columnar/columnar_row_ref.h"
+#include "paimon/common/table/special_fields.h"
+#include "paimon/common/types/data_field.h"
+#include "paimon/common/types/row_kind.h"
+#include "paimon/common/utils/arrow/arrow_utils.h"
+#include "paimon/common/utils/arrow/mem_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+#include "paimon/common/utils/fields_comparator.h"
+#include "paimon/common/utils/scope_guard.h"
+#include "paimon/core/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,
+ const std::shared_ptr<arrow::MemoryPool>& arrow_pool) {
+ if (sealed_offsets.begin < 0 || sealed_offsets.end <
sealed_offsets.begin) {
+ return Status::Invalid("PK real-time store returned an invalid
sealed offset range");
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ std::shared_ptr<arrow::Buffer> seen_offsets,
+ arrow::AllocateEmptyBitmap(sealed_offsets.Count(),
arrow_pool.get()));
+ return std::shared_ptr<RealtimeOffsetCoverage>(new
RealtimeOffsetCoverage(
+ sealed_offsets, reader_count, std::move(seen_offsets),
arrow_pool));
+ }
+
+ Status Add(const arrow::Int64Array& offsets) {
+ std::lock_guard<std::mutex> lock(mutex_);
+ for (int64_t row = 0; row < offsets.length(); ++row) {
+ const int64_t offset = offsets.Value(row);
+ if (offset < sealed_offsets_.begin || offset >=
sealed_offsets_.end) {
+ return Status::Invalid(
+ "PK real-time store commit reader offset is outside the
sealed range");
+ }
+ const int64_t index = offset - sealed_offsets_.begin;
+ if (arrow::bit_util::GetBit(seen_offsets_->data(), index)) {
+ return Status::Invalid(
+ "PK real-time store commit readers contain duplicate
REALTIME_OFFSET");
+ }
+ arrow::bit_util::SetBit(seen_offsets_->mutable_data(), index);
+ ++seen_count_;
+ }
+ return Status::OK();
+ }
+
+ Status FinishReader() {
+ std::lock_guard<std::mutex> lock(mutex_);
+ ++finished_reader_count_;
+ if (finished_reader_count_ == reader_count_ && seen_count_ !=
sealed_offsets_.Count()) {
+ return Status::Invalid(
+ "PK real-time store commit readers did not cover the sealed
range");
+ }
+ return Status::OK();
+ }
+
+ private:
+ RealtimeOffsetCoverage(const OffsetRange& sealed_offsets, size_t
reader_count,
+ std::shared_ptr<arrow::Buffer> seen_offsets,
+ const std::shared_ptr<arrow::MemoryPool>&
arrow_pool)
+ : sealed_offsets_(sealed_offsets),
+ reader_count_(reader_count),
+ arrow_pool_(arrow_pool),
+ seen_offsets_(std::move(seen_offsets)) {}
+
+ OffsetRange sealed_offsets_;
+ size_t reader_count_;
+ std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+ std::shared_ptr<arrow::Buffer> seen_offsets_;
+ int64_t seen_count_ = 0;
+ size_t finished_reader_count_ = 0;
+ std::mutex mutex_;
+};
+
+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;
+}
+
+Result<std::shared_ptr<arrow::StructArray>> ApplyOffsetFilter(
+ const std::shared_ptr<arrow::StructArray>& data_batch,
+ const std::shared_ptr<arrow::NumericArray<arrow::Int64Type>>& offset_array,
+ const std::optional<OffsetRange>& visible_offsets, arrow::MemoryPool*
arrow_pool) {
+ if (!visible_offsets.has_value()) {
+ return data_batch;
+ }
+
+ arrow::BooleanBuilder filter_builder(arrow_pool);
+
PAIMON_RETURN_NOT_OK_FROM_ARROW(filter_builder.Reserve(offset_array->length()));
+ int64_t visible_row_count = 0;
+ for (int64_t i = 0; i < offset_array->length(); ++i) {
+ int64_t offset = offset_array->Value(i);
+ bool visible = offset >= visible_offsets->begin && offset <
visible_offsets->end;
+ filter_builder.UnsafeAppend(visible);
+ visible_row_count += visible;
+ }
+ if (visible_row_count == 0) {
+ return std::shared_ptr<arrow::StructArray>();
+ }
+ if (visible_row_count == data_batch->length()) {
+ return data_batch;
+ }
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> filter,
+ filter_builder.Finish());
+ arrow::compute::ExecContext exec_context(arrow_pool);
+ PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+ arrow::Datum filtered,
+ arrow::compute::Filter(data_batch, filter,
arrow::compute::FilterOptions::Defaults(),
+ &exec_context));
+ return checked_pointer_cast<arrow::StructArray>(filtered.make_array());
+}
+
+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_->row_kind_array_->length();
+ }
+
+ Result<KeyValue> Next() override {
+ if (cursor_ >= reader_->row_kind_array_->length()) {
+ return Status::Invalid("No more prepared key values in current
iterator");
+ }
+ std::shared_ptr<InternalRow> key =
+ std::make_shared<ColumnarRowRef>(reader_->key_ctx_, cursor_);
+ auto value = std::make_unique<ColumnarRowRef>(reader_->value_ctx_,
cursor_);
+ PAIMON_ASSIGN_OR_RAISE(
+ const RowKind* row_kind,
+
RowKind::FromByteValue(reader_->row_kind_array_->Value(cursor_)));
+ int64_t sequence_number =
reader_->sequence_number_array_->Value(cursor_);
+ ++cursor_;
+ return KeyValue(row_kind, sequence_number,
KeyValue::UNKNOWN_LEVEL, std::move(key),
+ std::move(value));
+ }
+
+ private:
+ PreparedKeyValueReader* reader_;
+ int64_t cursor_ = 0;
+ };
+
+ Result<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(
+ arrow::schema(data_batch->type()->fields()), key_schema_));
+ PAIMON_ASSIGN_OR_RAISE(
+ arrow_array,
+ AlignArrayByPaimonIds(data_batch,
arrow::struct_(prepared_schema_->fields()),
+ arrow_pool_.get()));
+ data_batch =
checked_pointer_cast<arrow::StructArray>(arrow_array);
+ }
+ PAIMON_RETURN_NOT_OK(ValidatePreparedBatch(data_batch));
+ PAIMON_RETURN_NOT_OK(ValidateOrdering(data_batch));
+
+ std::shared_ptr<arrow::NumericArray<arrow::Int64Type>>
offset_array =
+ checked_pointer_cast<arrow::NumericArray<arrow::Int64Type>>(
+ data_batch->field(kRealtimeOffsetIndex));
+ if (offset_coverage_) {
+ PAIMON_RETURN_NOT_OK(offset_coverage_->Add(*offset_array));
+ }
+ PAIMON_ASSIGN_OR_RAISE(
+ data_batch,
+ ApplyOffsetFilter(data_batch, offset_array, visible_offsets_,
arrow_pool_.get()));
+ if (!data_batch) {
+ continue;
+ }
+
+ row_kind_array_ =
checked_pointer_cast<arrow::NumericArray<arrow::Int8Type>>(
+ data_batch->field(kValueKindIndex));
+ sequence_number_array_ =
checked_pointer_cast<arrow::NumericArray<arrow::Int64Type>>(
+ data_batch->field(kSequenceNumberIndex));
+ PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector key_fields,
+ ProjectFieldsByPaimonIds(data_batch,
prepared_schema_,
+ key_schema_,
arrow_pool_.get()));
+ PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector value_fields,
+ ProjectFieldsByPaimonIds(data_batch,
prepared_schema_,
+ value_schema_,
arrow_pool_.get()));
+ key_ctx_ = std::make_shared<ColumnarBatchContext>(key_fields,
pool_);
+ value_ctx_ = std::make_shared<ColumnarBatchContext>(value_fields,
pool_);
+ ArrowUtils::TraverseArray(data_batch);
+ return std::make_unique<Iterator>(this);
+ }
+ }
+
+ Status ValidatePreparedBatch(const std::shared_ptr<arrow::StructArray>&
data_batch) const {
+ if (data_batch->num_fields() != prepared_schema_->num_fields()) {
+ return Status::Invalid(fmt::format(
+ "prepared batch field count {} does not match prepared schema
field count {}",
+ data_batch->num_fields(), prepared_schema_->num_fields()));
+ }
+ const arrow::FieldVector& batch_fields = data_batch->type()->fields();
+ for (int32_t i = 0; i < data_batch->num_fields(); ++i) {
+ if (!batch_fields[i]->Equals(prepared_schema_->field(i), true)) {
+ return Status::Invalid(fmt::format(
+ "prepared batch field {} does not match declared prepared
schema", i));
+ }
+ }
+ if (!data_batch->field(kValueKindIndex) ||
+ data_batch->field(kValueKindIndex)->type_id() !=
arrow::Type::INT8) {
+ return Status::Invalid("cannot cast VALUE_KIND column to int8
arrow array");
+ }
+ if (!data_batch->field(kSequenceNumberIndex) ||
+ data_batch->field(kSequenceNumberIndex)->type_id() !=
arrow::Type::INT64) {
+ return Status::Invalid("cannot cast SEQUENCE_NUMBER column to
int64 arrow array");
+ }
+ if (!data_batch->field(kRealtimeOffsetIndex) ||
+ data_batch->field(kRealtimeOffsetIndex)->type_id() !=
arrow::Type::INT64) {
+ return Status::Invalid("cannot cast REALTIME_OFFSET column to
int64 arrow array");
+ }
+ if (data_batch->field(kValueKindIndex)->null_count() != 0 ||
+ data_batch->field(kSequenceNumberIndex)->null_count() != 0 ||
+ data_batch->field(kRealtimeOffsetIndex)->null_count() != 0) {
+ return Status::Invalid("prepared transport columns must not
contain nulls");
+ }
+ return Status::OK();
+ }
+
+ Status ValidateOrdering(const std::shared_ptr<arrow::StructArray>&
data_batch) {
Review Comment:
Thanks. The ordering scan was originally added to validate custom store
output before MOR. However, sorted output is already part of the RealtimeStore
reader contract, and validating every adjacent row duplicated key projection
and comparison performed by the downstream merge reader. Commit
b2827df103f32c533b7595205fa9bcd7655b18c6 removes that hot-path validation and
leaves sorting responsibility at the framework/store contract boundary.
--
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]
