HaHaJeff commented on code in PR #224: URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3852188106
########## 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 at current PR head `a7bc03b7b23c0302ad77d0e7312c0fb4da8a334c`. `PrimaryKeyRealtimeStore::CreateQueryReaders` now aligns each stored batch to `RealtimeQueryContext::read_schema` through `NestedProjectionUtils::AlignArrayToReadType`, matching the append-store query path. `PreparedKeyValueReader` only converts the already projected transport schema; the earlier framework-side recursive schema-reconciliation path has been removed. Schema evolution remains unsupported in V1 and requires recreating the `RealtimeContext`/store. `TestPkNestedProjectionAcrossDiskAndMemory` covers disk, sealed memory, and active memory with nested projection. ########## 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) { Review Comment: The earlier count/min/max implementation was superseded by the exact commit-coverage requirement. At current PR head `a7bc03b7b23c0302ad77d0e7312c0fb4da8a334c`, commit readers share a `RealtimeOffsetCoverage` backed by `RoaringBitmap64`: each offset is checked for range and duplication, and complete sealed-range cardinality is verified after all readers reach EOF. Query readers do not use this coverage object. This keeps runtime validation because PK-sorted readers do not preserve offset monotonicity, so count/min/max cannot detect a duplicate-plus-hole case. Integration tests cover missing, duplicate, and out-of-range offsets. -- 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]
