HaHaJeff commented on code in PR #224: URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3828375726
########## 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: Fixed in 5f74e463, with schema-evolution and field-ID follow-ups in 43afe022 and 9d2d19c2. Stored PK batches are aligned to the requested nested type before constructing the KeyValue reader, and the added unit/integration coverage verifies payload<b> across memory and disk. -- 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]
