HaHaJeff commented on code in PR #224: URL: https://github.com/apache/paimon-cpp/pull/224#discussion_r3873133090
########## src/paimon/core/realtime/realtime_primary_key_writer.cpp: ########## @@ -0,0 +1,279 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/realtime/realtime_primary_key_writer.h" + +#include <limits> +#include <optional> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/api.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/deduplicate_merge_function.h" +#include "paimon/core/mergetree/compact/reducer_merge_function_wrapper.h" +#include "paimon/core/mergetree/merge_tree_writer.h" +#include "paimon/core/realtime/prepared_key_value_reader.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +Result<std::shared_ptr<arrow::StructArray>> PrepareBatch( + std::unique_ptr<RecordBatch>&& batch, const std::shared_ptr<arrow::Schema>& write_schema, + const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::vector<std::string>& trimmed_primary_keys, int64_t first_sequence_number, + int64_t first_offset, arrow::MemoryPool* arrow_pool) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> input, + arrow::ImportArray(batch->GetData(), arrow::struct_(write_schema->fields()))); + if (!input || input->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time write data is not a StructArray"); + } + std::shared_ptr<arrow::StructArray> values = checked_pointer_cast<arrow::StructArray>(input); + const int64_t count = values->length(); + arrow::Int8Builder kinds(arrow_pool); + arrow::Int64Builder sequences(arrow_pool); + arrow::Int64Builder offsets(arrow_pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Reserve(count)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Reserve(count)); + const std::vector<RecordBatch::RowKind>& row_kinds = batch->GetRowKind(); + for (int64_t row = 0; row < count; ++row) { + const RecordBatch::RowKind kind = + row_kinds.empty() ? RecordBatch::RowKind::INSERT : row_kinds[row]; + kinds.UnsafeAppend(static_cast<int8_t>(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr<arrow::Array> kind_array; + std::shared_ptr<arrow::Array> sequence_array; + std::shared_ptr<arrow::Array> offset_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(kinds.Finish(&kind_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(sequences.Finish(&sequence_array)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(offsets.Finish(&offset_array)); + arrow::ArrayVector columns = {std::move(kind_array), std::move(sequence_array), + std::move(offset_array)}; + columns.insert(columns.end(), values->fields().begin(), values->fields().end()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::StructArray> prepared, + arrow::StructArray::Make(std::move(columns), prepared_schema->fields())); + + std::vector<arrow::compute::SortKey> sort_keys; + sort_keys.reserve(trimmed_primary_keys.size() + 1); + for (const std::string& key : trimmed_primary_keys) { + sort_keys.emplace_back(key, arrow::compute::SortOrder::Ascending); + } + sort_keys.emplace_back(SpecialFields::SequenceNumber().Name(), + arrow::compute::SortOrder::Ascending); + arrow::compute::ExecContext context(arrow_pool); + arrow::compute::SortOptions options(sort_keys, arrow::compute::NullPlacement::AtStart); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum indices, + arrow::compute::SortIndices(arrow::Datum(prepared), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(prepared), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + return checked_pointer_cast<arrow::StructArray>(sorted.make_array()); +} + +} // namespace + +Result<std::shared_ptr<RealtimePrimaryKeyWriter>> RealtimePrimaryKeyWriter::Create( + const std::map<std::string, std::string>& partition, int32_t bucket, + const std::shared_ptr<arrow::Schema>& write_schema, + const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::vector<std::string>& trimmed_primary_keys, + const std::shared_ptr<FieldsComparator>& key_comparator, + const std::shared_ptr<RealtimeContextImpl>& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr<MergeTreeWriter>& merge_tree_writer, + const std::shared_ptr<MemoryPool>& memory_pool) { + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits<int64_t>::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr<arrow::Field> field = write_schema->GetFieldByName(key); + if (!field) { + return Status::Invalid("PK field is missing from write schema: ", key); + } + key_fields.push_back(std::move(field)); + } + const RealtimePartitionBucket partition_bucket(partition, bucket); + PAIMON_ASSIGN_OR_RAISE(int64_t initial_max_sequence_number, + realtime_context->AdvanceMaterializedMaxSequenceNumber( + partition_bucket, restored_max_sequence_number)); + return std::shared_ptr<RealtimePrimaryKeyWriter>(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, + prepared_schema, arrow::schema(std::move(key_fields)), trimmed_primary_keys, key_comparator, + store_state.initial_offset, initial_max_sequence_number, memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr<RealtimeStore>& realtime_store, + const std::shared_ptr<MergeTreeWriter>& merge_tree_writer, + const std::shared_ptr<RealtimeContextImpl>& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr<arrow::Schema>& write_schema, + const std::shared_ptr<arrow::Schema>& prepared_schema, + const std::shared_ptr<arrow::Schema>& key_schema, + const std::vector<std::string>& trimmed_primary_keys, + const std::shared_ptr<FieldsComparator>& key_comparator, int64_t next_offset, + int64_t last_sequence_number, const std::shared_ptr<MemoryPool>& memory_pool) + : memory_pool_(memory_pool), + arrow_pool_(GetArrowPool(memory_pool)), + realtime_store_(realtime_store), + merge_tree_writer_(merge_tree_writer), + realtime_context_(realtime_context), + partition_bucket_(partition_bucket), + write_schema_(write_schema), + prepared_schema_(prepared_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr<RecordBatch>&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("PK real-time write batch is null"); + } + const int64_t count = batch->GetData()->length; + if (count == 0) { + return Status::OK(); + } + const std::vector<RecordBatch::RowKind>& row_kinds = batch->GetRowKind(); + if (!row_kinds.empty() && static_cast<int64_t>(row_kinds.size()) != count) { + return Status::Invalid("PK real-time row-kind count does not match batch row count"); + } + for (RecordBatch::RowKind row_kind : row_kinds) { + PAIMON_ASSIGN_OR_RAISE(const RowKind* validated, + RowKind::FromByteValue(static_cast<int8_t>(row_kind))); + static_cast<void>(validated); + } + std::lock_guard<std::mutex> lock(realtime_store_mutex_); + if (count > std::numeric_limits<int64_t>::max() - next_offset_) { + return Status::Invalid("real-time offset range exceeds INT64_MAX"); + } + // Reserve INT64_MAX as the exhausted sequence-number sentinel. + if (last_sequence_number_ >= std::numeric_limits<int64_t>::max() - count) { + return Status::Invalid("PK sequence range exceeds INT64_MAX"); + } + const int64_t first_sequence = last_sequence_number_ + 1; + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<arrow::StructArray> prepared, + PrepareBatch(std::move(batch), write_schema_, prepared_schema_, trimmed_primary_keys_, + first_sequence, next_offset_, arrow_pool_.get())); + auto output = std::make_unique<ArrowArray>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*prepared, output.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> prepared_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(prepared_batch), OffsetRange(next_offset_, next_offset_ + count)})); + next_offset_ += count; + last_sequence_number_ += count; + PAIMON_RETURN_NOT_OK(realtime_context_->AdvanceMaterializedMaxSequenceNumber( + partition_bucket_, last_sequence_number_)); + return Status::OK(); +} + +Result<CommitIncrement> RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard<std::mutex> prepare_lock(prepare_mutex_); + std::optional<std::shared_ptr<RealtimeSegmentHandle>> segment; + { + std::lock_guard<std::mutex> store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<RealtimeSegmentHandle>> sealed, + realtime_store_->SealForCommit()); + segment = std::move(sealed); + } + if (segment && !segment.value()) { + return Status::Invalid("PK real-time store sealed a null segment"); + } + std::optional<OffsetRange> sealed_range; + if (segment) { + sealed_range = segment.value()->GetOffsetRange(); + if (sealed_range->begin < 0 || sealed_range->end < sealed_range->begin) { + return Status::Invalid("PK real-time store returned an invalid sealed offset range"); + } + PAIMON_RETURN_NOT_OK(FlushSegment(segment.value(), sealed_range.value())); + } + PAIMON_ASSIGN_OR_RAISE(CommitIncrement increment, + merge_tree_writer_->PrepareCommit(wait_compaction)); + if (segment) { + increment.SetRealtimeOffsetRange(sealed_range.value()); + } + return increment; +} + +Status RealtimePrimaryKeyWriter::FlushSegment(const std::shared_ptr<RealtimeSegmentHandle>& segment, + const OffsetRange& sealed_offsets) { + PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<BatchReader>> readers, + realtime_store_->CreateCommitReaders(segment)); + PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<KeyValueRecordReader>> prepared_readers, + PreparedKeyValueReaderFactory::CreateForCommit( + std::move(readers), prepared_schema_, sealed_offsets, key_schema_, + write_schema_, memory_pool_)); + std::vector<std::unique_ptr<KeyValueRecordReader>> sorted_readers; + sorted_readers.reserve(prepared_readers.size()); + for (std::unique_ptr<KeyValueRecordReader>& prepared_reader : prepared_readers) { + auto merge_function = std::make_unique<DeduplicateMergeFunction>(/*ignore_delete=*/false); + sorted_readers.push_back(std::make_unique<MergedKeyValueRecordReader>( Review Comment: The writer now rejects non-deduplicate engines during construction, while the flush path creates merge functions through `PrimaryKeyTableUtils::CreateMergeFunction`. Supporting another engine later should only require relaxing the construction-time validation. Implemented in `982eb59256ebf14d943e202950b35ad30b8fc671`. ########## src/paimon/core/operation/merge_file_split_read.cpp: ########## @@ -453,38 +611,90 @@ Result<std::unique_ptr<BatchReader>> MergeFileSplitRead::CreateReaderForSection( } else { predicate = context_->GetPredicate(); } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<SortMergeReader> sort_merge_reader, - CreateSortMergeReaderForSection(section, partition, dv_factory, - predicate, data_file_path_factory, - /*drop_delete=*/!force_keep_delete_)); - // KeyValueProjectionReader converts KeyValue objects to arrow array according to projection - if (!context_->EnableMultiThreadRowToBatch()) { - return KeyValueProjectionReader::Create(std::move(sort_merge_reader), raw_read_schema_, - projection_, options_.GetReadBatchSize(), pool_); - } - int32_t thread_number = context_->GetRowToBatchThreadNumber(); - assert(thread_number > 0); - return std::make_unique<AsyncKeyValueProjectionReader>( - std::move(sort_merge_reader), raw_read_schema_, projection_, options_.GetReadBatchSize(), - thread_number, pool_); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<SortMergeReader> sort_merge_reader, + CreateSortMergeReaderForSection(section, partition, dv_factory, predicate, + data_file_path_factory, /*drop_delete=*/false)); + return CreateProjectedReader(std::move(sort_merge_reader), /*predicate=*/nullptr, + /*complete_row_kind=*/false); } -Result<std::unique_ptr<SortMergeReader>> MergeFileSplitRead::CreateSortMergeReaderForSection( +Status MergeFileSplitRead::CreateDiskSections( + const std::vector<std::shared_ptr<DataFileMeta>>& data_files, + const std::vector<std::optional<DeletionFile>>& deletion_files, + DeletionVector::Factory* dv_factory, std::vector<std::vector<SortedRun>>* sections) const { + *dv_factory = DeletionVector::CreateFactory( + options_.GetFileSystem(), DeletionVector::CreateDeletionFileMap(data_files, deletion_files), + pool_); + *sections = IntervalPartition(data_files, key_comparator_).Partition(); + return Status::OK(); +} + +Result<std::vector<std::unique_ptr<KeyValueRecordReader>>> +MergeFileSplitRead::CreateRecordReadersForSection( const std::vector<SortedRun>& section, const BinaryRow& partition, DeletionVector::Factory dv_factory, const std::shared_ptr<Predicate>& predicate, - const std::shared_ptr<DataFilePathFactory>& data_file_path_factory, bool drop_delete) { - // with overlap in one section + const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const { Review Comment: Restored the original comment above the corresponding `KeyValueProjectionReader` construction in `982eb59256ebf14d943e202950b35ad30b8fc671`. ########## src/paimon/core/realtime/realtime_primary_key_reader.h: ########## @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <memory> +#include <vector> + +#include "arrow/type_fwd.h" +#include "paimon/core/io/key_value_record_reader.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/result.h" + +namespace paimon { +class BatchReader; +class MemoryPool; + +class PreparedKeyValueReaderFactory { + public: Review Comment: `Prepared*` has been replaced with explicit PK realtime terminology. The reader, layout helper, files, tests, and related variables now use `RealtimePrimaryKey*` names. Updated in `982eb59256ebf14d943e202950b35ad30b8fc671`. -- 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]
