wangyong9999 commented on code in PR #224:
URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3826984194


##########
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:
   Only the watermark returned to the new `MergeTreeWriter` is advanced here. 
When `stores_` already contains this partition-bucket, 
`PrimaryKeyRealtimeStore::next_sequence_number_` remains on the old watermark. 
If the latest snapshot has advanced from sequence 4 to 10, the next in-memory 
mutation can get 5 while the same row gets 11 when flushed; merge-on-read then 
lets the disk row at 10 hide the newer memory row. Advance or reject the reused 
store atomically, or use one shared sequence allocator for both paths.



##########
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:
   This maps nested projections only at the top level. With stored 
`payload<a,b>` and requested `payload<b>`, `KeyValueProjectionReader` builds 
the pruned struct but reads child 0 from the full stored row, returning `a` as 
`b` when their types match. Align each stored batch to the requested nested 
type—as `ArrowRealtimeStore` does with `AlignArrayToReadType`—before building 
the `KeyValue` reader, and add a memory-plus-disk nested-projection test.



-- 
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]

Reply via email to