diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index bf57ce9f..051eba32 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -390,6 +390,7 @@ set(PAIMON_CORE_SRCS core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp + core/realtime/realtime_primary_key_writer.cpp core/postpone/postpone_bucket_writer.cpp core/schema/arrow_schema_validator.cpp core/schema/schema_manager.cpp diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp new file mode 100644 index 00000000..61cd1400 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -0,0 +1,293 @@ +/* + * 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 +#include +#include +#include + +#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/core_options.h" +#include "paimon/core/io/merged_key_value_record_reader.h" +#include "paimon/core/mergetree/compact/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/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_primary_key_reader.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/core/utils/primary_key_table_utils.h" +#include "paimon/macros.h" + +namespace paimon { + +namespace { + +Result> CreateRealtimePrimaryKeyTransportBatch( + std::unique_ptr&& batch, const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& 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 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 values = checked_pointer_cast(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& 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(kind)); + sequences.UnsafeAppend(first_sequence_number + row); + offsets.UnsafeAppend(first_offset + row); + } + std::shared_ptr kind_array; + std::shared_ptr sequence_array; + std::shared_ptr 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 transport, + arrow::StructArray::Make(std::move(columns), transport_schema->fields())); + + std::vector 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(transport), options, &context)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum sorted, + arrow::compute::Take(arrow::Datum(transport), indices, + arrow::compute::TakeOptions::NoBoundsCheck(), &context)); + return checked_pointer_cast(sorted.make_array()); +} + +} // namespace + +Result> RealtimePrimaryKeyWriter::Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool) { + if (restored_max_sequence_number < -1 || + restored_max_sequence_number == std::numeric_limits::max()) { + return Status::Invalid("PK restored sequence number is invalid"); + } + if (options.GetMergeEngine() != MergeEngine::DEDUPLICATE) { + return Status::NotImplemented("PK realtime supports only the DEDUPLICATE merge engine"); + } + PAIMON_RETURN_NOT_OK(RealtimePrimaryKeyLayout::ValidateSchema(transport_schema)); + arrow::FieldVector key_fields; + key_fields.reserve(trimmed_primary_keys.size()); + for (const std::string& key : trimmed_primary_keys) { + std::shared_ptr 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(new RealtimePrimaryKeyWriter( + store_state.store, merge_tree_writer, realtime_context, partition_bucket, write_schema, + transport_schema, arrow::schema(std::move(key_fields)), trimmed_primary_keys, + key_comparator, options, store_state.initial_offset, initial_max_sequence_number, + memory_pool)); +} + +RealtimePrimaryKeyWriter::RealtimePrimaryKeyWriter( + const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + int64_t next_offset, int64_t last_sequence_number, + const std::shared_ptr& 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), + transport_schema_(transport_schema), + key_schema_(key_schema), + trimmed_primary_keys_(trimmed_primary_keys), + key_comparator_(key_comparator), + options_(options), + next_offset_(next_offset), + last_sequence_number_(last_sequence_number) {} + +Status RealtimePrimaryKeyWriter::Write(std::unique_ptr&& 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& row_kinds = batch->GetRowKind(); + if (!row_kinds.empty() && static_cast(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(row_kind))); + static_cast(validated); + } + std::lock_guard lock(realtime_store_mutex_); + if (count > std::numeric_limits::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::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 transport, + CreateRealtimePrimaryKeyTransportBatch(std::move(batch), write_schema_, transport_schema_, + trimmed_primary_keys_, first_sequence, next_offset_, + arrow_pool_.get())); + auto output = std::make_unique(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*transport, output.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(output.get(), arrow_pool_)); + RecordBatchBuilder builder(output.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr transport_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK(realtime_store_->Write(RealtimeWriteBatch{ + std::move(transport_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 RealtimePrimaryKeyWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard prepare_lock(prepare_mutex_); + std::optional> segment; + { + std::lock_guard store_lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(std::optional> 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 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& segment, + const OffsetRange& sealed_offsets) { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + realtime_store_->CreateCommitReaders(segment)); + PAIMON_ASSIGN_OR_RAISE( + std::vector> realtime_primary_key_readers, + RealtimePrimaryKeyReaderFactory::CreateForCommit(std::move(readers), transport_schema_, + sealed_offsets, key_schema_, write_schema_, + memory_pool_)); + std::vector> sorted_readers; + sorted_readers.reserve(realtime_primary_key_readers.size()); + for (std::unique_ptr& realtime_primary_key_reader : + realtime_primary_key_readers) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr merge_function, + PrimaryKeyTableUtils::CreateMergeFunction( + write_schema_, trimmed_primary_keys_, options_, memory_pool_)); + sorted_readers.push_back(std::make_unique( + std::move(realtime_primary_key_reader), key_comparator_, + std::make_shared(std::move(merge_function)))); + } + return merge_tree_writer_->WriteSortedReadersToFiles(std::move(sorted_readers)); +} + +Status RealtimePrimaryKeyWriter::Compact(bool) { + return Status::Invalid("PK real-time write does not support explicit compaction"); +} +uint64_t RealtimePrimaryKeyWriter::GetMemoryUsage() const { + return realtime_store_->GetMemoryUsage(); +} +Status RealtimePrimaryKeyWriter::FlushMemory() { + return Status::OK(); +} +Result RealtimePrimaryKeyWriter::CompactNotCompleted() { + return merge_tree_writer_->CompactNotCompleted(); +} +Status RealtimePrimaryKeyWriter::Sync() { + return merge_tree_writer_->Sync(); +} +Status RealtimePrimaryKeyWriter::Close() { + return merge_tree_writer_->Close(); +} +std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { + return merge_tree_writer_->GetMetrics(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h new file mode 100644 index 00000000..cdd3d889 --- /dev/null +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -0,0 +1,105 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include "paimon/core/core_options.h" +#include "paimon/core/utils/batch_writer.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/realtime/realtime_store.h" + +namespace arrow { +class MemoryPool; +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; +class MergeTreeWriter; +class FieldsComparator; +class RealtimeContextImpl; +struct RealtimeStoreState; + +class RealtimePrimaryKeyWriter final : public BatchWriter { + public: + static Result> Create( + const std::map& partition, int32_t bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, const CoreOptions& options, + const std::shared_ptr& realtime_context, + const RealtimeStoreState& store_state, int64_t restored_max_sequence_number, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& memory_pool); + + Status Write(std::unique_ptr&& batch) override; + Result PrepareCommit(bool wait_compaction) override; + Status Compact(bool full_compaction) override; + uint64_t GetMemoryUsage() const override; + Status FlushMemory() override; + Result CompactNotCompleted() override; + Status Sync() override; + Status Close() override; + std::shared_ptr GetMetrics() const override; + + private: + RealtimePrimaryKeyWriter(const std::shared_ptr& realtime_store, + const std::shared_ptr& merge_tree_writer, + const std::shared_ptr& realtime_context, + const RealtimePartitionBucket& partition_bucket, + const std::shared_ptr& write_schema, + const std::shared_ptr& transport_schema, + const std::shared_ptr& key_schema, + const std::vector& trimmed_primary_keys, + const std::shared_ptr& key_comparator, + const CoreOptions& options, int64_t next_offset, + int64_t last_sequence_number, + const std::shared_ptr& memory_pool); + + Status FlushSegment(const std::shared_ptr& segment, + const OffsetRange& sealed_offsets); + + std::shared_ptr memory_pool_; + std::shared_ptr arrow_pool_; + std::shared_ptr realtime_store_; + std::shared_ptr merge_tree_writer_; + std::shared_ptr realtime_context_; + RealtimePartitionBucket partition_bucket_; + std::shared_ptr write_schema_; + std::shared_ptr transport_schema_; + std::shared_ptr key_schema_; + std::vector trimmed_primary_keys_; + std::shared_ptr key_comparator_; + CoreOptions options_; + int64_t next_offset_; + int64_t last_sequence_number_; + std::mutex realtime_store_mutex_; + std::mutex prepare_mutex_; +}; + +} // namespace paimon