zjw1111 commented on code in PR #393: URL: https://github.com/apache/paimon-cpp/pull/393#discussion_r4280473308
########## src/paimon/core/realtime/arrow_deduplicate_realtime_store.cpp: ########## @@ -0,0 +1,542 @@ +/* + * 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/arrow_deduplicate_realtime_store.h" + +#include <algorithm> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "arrow/compute/api.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/table/special_fields.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/realtime/realtime_utils.h" +#include "paimon/macros.h" +#include "paimon/metrics.h" +#include "paimon/record_batch.h" + +namespace paimon { +namespace { + +Result<std::shared_ptr<arrow::StructArray>> ProjectByName( + const std::shared_ptr<arrow::StructArray>& source, + const std::shared_ptr<arrow::Schema>& target_schema) { + arrow::ArrayVector fields; + fields.reserve(target_schema->num_fields()); + for (const std::shared_ptr<arrow::Field>& field : target_schema->fields()) { + std::shared_ptr<arrow::Array> array = source->GetFieldByName(field->name()); + if (!array || !array->type()->Equals(field->type())) { + return Status::Invalid("real-time deduplicate reader cannot project field: ", + field->name()); + } + fields.push_back(std::move(array)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> result, + arrow::StructArray::Make(fields, target_schema->fields())); + return result; +} + +Result<BatchReader::ReadBatch> ExportBatch(const std::shared_ptr<arrow::StructArray>& data, + const std::shared_ptr<arrow::MemoryPool>& pool) { + auto array = std::make_unique<ArrowArray>(); + auto schema = std::make_unique<ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*data, array.get(), schema.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(array.get(), schema.get(), pool)); + return std::make_pair(std::move(array), std::move(schema)); +} + +Result<std::shared_ptr<arrow::StructArray>> FilterDeletedOffsets( + const std::shared_ptr<arrow::StructArray>& source, + const std::shared_ptr<const RoaringBitmap64>& deleted_offsets, + const std::shared_ptr<arrow::MemoryPool>& pool) { + if (!deleted_offsets || deleted_offsets->IsEmpty() || source->length() == 0) { + return source; + } + std::shared_ptr<arrow::Array> offset_array = + source->GetFieldByName(SpecialFields::RealtimeOffset().Name()); + if (!offset_array || offset_array->type_id() != arrow::Type::INT64 || + offset_array->null_count() != 0) { + return Status::Invalid( + "real-time deduplicate data must contain non-null int64 _REALTIME_OFFSET"); + } + std::shared_ptr<arrow::Int64Array> offsets = + checked_pointer_cast<arrow::Int64Array>(offset_array); + arrow::BooleanBuilder selection(pool.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(selection.Reserve(source->length())); + // TODO(xinyu.lxy): Merge-scan the ordered offsets with a RoaringBitmap64 iterator instead of + // calling Contains for every row. + for (int64_t row = 0; row < source->length(); ++row) { + selection.UnsafeAppend(!deleted_offsets->Contains(offsets->Value(row))); + } + std::shared_ptr<arrow::Array> selection_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(selection.Finish(&selection_array)); + arrow::compute::ExecContext exec_context(pool.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum filtered, + arrow::compute::Filter(source, selection_array, arrow::compute::FilterOptions::Defaults(), + &exec_context)); + return checked_pointer_cast<arrow::StructArray>(filtered.make_array()); +} + +} // namespace + +class ArrowDeduplicateRealtimeStore::SegmentHandle final : public RealtimeSegmentHandle { + public: + SegmentHandle(std::shared_ptr<RealtimeSegmentHandle> delegate, OffsetRange offset_range, + int64_t row_count, std::shared_ptr<const RoaringBitmap64> offset_deletions) + : delegate_(std::move(delegate)), + offset_range_(offset_range), + row_count_(row_count), + offset_deletions_(std::move(offset_deletions)) {} + + OffsetRange GetOffsetRange() const override { + return offset_range_; + } + + int64_t GetRowCount() const override { + return row_count_; + } + + const std::shared_ptr<RealtimeSegmentHandle>& Delegate() const { + return delegate_; + } + + const std::shared_ptr<const RoaringBitmap64>& OffsetDeletions() const { + return offset_deletions_; + } + + private: + std::shared_ptr<RealtimeSegmentHandle> delegate_; + OffsetRange offset_range_; + int64_t row_count_; + std::shared_ptr<const RoaringBitmap64> offset_deletions_; +}; + +class ArrowDeduplicateRealtimeStore::ReadView final : public RealtimeReadView { + public: + ReadView(std::shared_ptr<RealtimeReadView> delegate, std::optional<OffsetRange> offset_range, + std::shared_ptr<const RoaringBitmap64> offset_deletions) + : delegate_(std::move(delegate)), + offset_range_(offset_range), + offset_deletions_(std::move(offset_deletions)) {} + + std::optional<OffsetRange> GetOffsetRange() const override { + return offset_range_; + } + + const std::shared_ptr<RealtimeReadView>& Delegate() const { + return delegate_; + } + + const std::shared_ptr<const RoaringBitmap64>& OffsetDeletions() const { + return offset_deletions_; + } + + private: + std::shared_ptr<RealtimeReadView> delegate_; + std::optional<OffsetRange> offset_range_; + std::shared_ptr<const RoaringBitmap64> offset_deletions_; +}; + +class ArrowDeduplicateRealtimeStore::OffsetFilteringBatchReader final : public BatchReader { + public: + OffsetFilteringBatchReader(std::unique_ptr<BatchReader> inner, + std::shared_ptr<arrow::Schema> output_schema, + std::shared_ptr<const RoaringBitmap64> offset_deletions, + std::shared_ptr<arrow::MemoryPool> arrow_pool) + : inner_(std::move(inner)), + output_schema_(std::move(output_schema)), + offset_deletions_(std::move(offset_deletions)), + arrow_pool_(std::move(arrow_pool)) {} + + // TODO(xinyu.lxy): Override NextBatchWithBitmap to compose the inner selection with offset + // deletions and avoid materializing an intermediate filtered array. + Result<ReadBatch> NextBatch() override { + while (true) { + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, inner_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> imported, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (!imported || imported->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("real-time deduplicate reader returned a non-StructArray"); + } + std::shared_ptr<arrow::StructArray> source = + checked_pointer_cast<arrow::StructArray>(imported); + PAIMON_ASSIGN_OR_RAISE(source, + FilterDeletedOffsets(source, offset_deletions_, arrow_pool_)); + if (source->length() == 0) { + continue; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::StructArray> projected, + ProjectByName(source, output_schema_)); + return ExportBatch(projected, arrow_pool_); + } + } + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return inner_ ? inner_->GetReaderMetrics() : std::make_shared<MetricsImpl>(); + } + + void Close() override { + if (inner_) { + inner_->Close(); + inner_.reset(); + } + } + + private: + std::unique_ptr<BatchReader> inner_; + std::shared_ptr<arrow::Schema> output_schema_; + std::shared_ptr<const RoaringBitmap64> offset_deletions_; + std::shared_ptr<arrow::MemoryPool> arrow_pool_; +}; + +Result<std::shared_ptr<ArrowDeduplicateRealtimeStore>> ArrowDeduplicateRealtimeStore::Create( + const std::shared_ptr<arrow::Schema>& write_schema, + const std::vector<std::string>& deduplicate_key_fields, + const std::shared_ptr<RealtimeStore>& delegate, const std::shared_ptr<MemoryPool>& memory_pool, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (!write_schema || !delegate || !memory_pool || !arrow_pool) { + return Status::Invalid("real-time deduplicate store is missing a required dependency"); + } + PAIMON_ASSIGN_OR_RAISE(int32_t key_position, RealtimeUtils::GetDeduplicateKeyPosition( + write_schema, deduplicate_key_fields)); + PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(write_schema)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<MutableKeyOffsetIndex> building_key_index, + MutableKeyOffsetIndex::Create(write_schema->field(key_position), memory_pool)); + return std::shared_ptr<ArrowDeduplicateRealtimeStore>(new ArrowDeduplicateRealtimeStore( + write_schema, key_position, write_schema->field(key_position), + std::move(building_key_index), delegate, memory_pool, arrow_pool)); +} + +ArrowDeduplicateRealtimeStore::ArrowDeduplicateRealtimeStore( + const std::shared_ptr<arrow::Schema>& write_schema, int32_t key_field_position, + const std::shared_ptr<arrow::Field>& key_field, + std::shared_ptr<MutableKeyOffsetIndex> building_key_index, + const std::shared_ptr<RealtimeStore>& delegate, const std::shared_ptr<MemoryPool>& memory_pool, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) + : write_schema_(write_schema), + key_field_position_(key_field_position), + key_field_(key_field), + delegate_(delegate), + memory_pool_(memory_pool), + arrow_pool_(arrow_pool), + building_key_index_(std::move(building_key_index)), + pending_offset_deletions_(std::make_shared<RoaringBitmap64>()) {} + +Status ArrowDeduplicateRealtimeStore::Write(RealtimeWriteBatch&&) { + return Status::Invalid( + "deduplicate store Write requires key-offset lookup; use RealtimeDeduplicateWriter"); +} + +Result<std::shared_ptr<arrow::StructArray>> ArrowDeduplicateRealtimeStore::ExtractKeys( + const std::shared_ptr<arrow::StructArray>& data) const { + if (!data || !data->type()->Equals(arrow::struct_(write_schema_->fields()))) { + return Status::Invalid("deduplicate write batch does not match store schema"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::StructArray> keys, + arrow::StructArray::Make({data->field(key_field_position_)}, {key_field_})); + return keys; +} + +Result<std::shared_ptr<const KeyOffsetLookup>> +ArrowDeduplicateRealtimeStore::AcquireKeyOffsetLookup() { + std::lock_guard<std::mutex> lock(mutex_); + if (!building_key_index_) { + return Status::Invalid("deduplicate building key-offset index is not initialized"); + } + std::vector<std::shared_ptr<const KeyOffsetLookup>> lookups; + if (committed_key_lookup_) { + lookups.push_back(committed_key_lookup_); + } + lookups.reserve(lookups.size() + sealed_states_.size() + 1); + for (const SealedState& sealed : sealed_states_) { + lookups.push_back(sealed.key_lookup); + } + lookups.push_back(building_key_index_->Snapshot()); + return std::shared_ptr<const KeyOffsetLookup>(new CompositeKeyOffsetLookup(std::move(lookups))); +} + +Status ArrowDeduplicateRealtimeStore::EnsureMutableDeletionBitmap() { + if (!pending_offset_deletions_.unique()) { + pending_offset_deletions_ = std::make_shared<RoaringBitmap64>(*pending_offset_deletions_); + } + return Status::OK(); +} + +Status ArrowDeduplicateRealtimeStore::WriteAndLookup( + const std::shared_ptr<arrow::StructArray>& data, + const std::vector<RecordBatch::RowKind>& row_kinds, const OffsetRange& offset_range, + const RoaringBitmap64& previous_offsets) { + if (!data || data->length() <= 0 || static_cast<int64_t>(row_kinds.size()) != data->length()) { + return Status::Invalid("deduplicate write data and row kinds are not aligned"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::StructArray> keys, ExtractKeys(data)); + std::shared_ptr<arrow::Array> raw_offsets = + data->GetFieldByName(SpecialFields::RealtimeOffset().Name()); + if (!raw_offsets || raw_offsets->type_id() != arrow::Type::INT64 || + raw_offsets->null_count() != 0) { + return Status::Invalid("deduplicate write requires non-null int64 offsets"); + } + std::shared_ptr<arrow::Int64Array> offsets = + checked_pointer_cast<arrow::Int64Array>(raw_offsets); + + std::lock_guard<std::mutex> lock(mutex_); + if (!building_key_index_) { + return Status::Invalid("deduplicate building key-offset index is not initialized"); + } + if (building_offset_range_ && offset_range.begin < building_offset_range_->end) { + return Status::Invalid("real-time offset ranges must be ordered and non-overlapping"); + } + PAIMON_RETURN_NOT_OK(EnsureMutableDeletionBitmap()); + *pending_offset_deletions_ |= previous_offsets; + + arrow::BooleanBuilder add_selection(arrow_pool_.get()); + PAIMON_RETURN_NOT_OK_FROM_ARROW(add_selection.Reserve(data->length())); + int64_t add_count = 0; + for (int64_t row = 0; row < data->length(); ++row) { + std::shared_ptr<arrow::StructArray> one_key = + checked_pointer_cast<arrow::StructArray>(keys->Slice(row, 1)); + PAIMON_ASSIGN_OR_RAISE(RoaringBitmap64 earlier_in_memory, + building_key_index_->LookupOffsets(one_key)); + *pending_offset_deletions_ |= earlier_in_memory; + + const RecordBatch::RowKind row_kind = row_kinds[row]; + const bool is_add = row_kind == RecordBatch::RowKind::INSERT || + row_kind == RecordBatch::RowKind::UPDATE_AFTER; + const bool is_delete = row_kind == RecordBatch::RowKind::DELETE || + row_kind == RecordBatch::RowKind::UPDATE_BEFORE; + if (!is_add && !is_delete) { + return Status::Invalid("unknown row kind in real-time deduplicate write"); + } + add_selection.UnsafeAppend(is_add); + if (!is_add) { + PAIMON_RETURN_NOT_OK(building_key_index_->Erase(one_key)); + continue; + } + std::shared_ptr<arrow::Int64Array> one_offset = + checked_pointer_cast<arrow::Int64Array>(offsets->Slice(row, 1)); + PAIMON_RETURN_NOT_OK(building_key_index_->Add(one_key, one_offset)); + building_row_offsets_.Add(offsets->Value(row)); + ++add_count; + } + + if (add_count > 0) { + std::shared_ptr<arrow::Array> selection; + PAIMON_RETURN_NOT_OK_FROM_ARROW(add_selection.Finish(&selection)); + arrow::compute::ExecContext exec_context(arrow_pool_.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum filtered, + arrow::compute::Filter(data, selection, arrow::compute::FilterOptions::Defaults(), + &exec_context)); + std::shared_ptr<arrow::StructArray> additions = + checked_pointer_cast<arrow::StructArray>(filtered.make_array()); + auto c_array = std::make_unique<ArrowArray>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*additions, c_array.get())); + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), /*schema=*/nullptr, arrow_pool_)); + RecordBatchBuilder builder(c_array.get()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<RecordBatch> append_batch, builder.Finish()); + PAIMON_RETURN_NOT_OK( + delegate_->Write(RealtimeWriteBatch{std::move(append_batch), offset_range})); Review Comment: [P1] A failed delegate write leaves deduplication state that can delete a retry The pending deletion bitmap and building key/row-offset indexes are modified before this call. If the delegate returns an error without accepting the batch, those changes remain, while `RealtimeDeduplicateWriter::Write()` does not advance `next_offset_` and `AbstractFileStoreWrite::Write()` only propagates the error, leaving the context usable. For example, replacing committed `k@1` with `k@2` can fail after recording `k -> 2`. Retrying offset 2 then finds that phantom offset and adds 2 to the deletion bitmap, so the successful retry is filtered out. A subsequent prepare can persist deletion of the old row without a replacement. ########## src/paimon/core/realtime/realtime_deduplicate_writer.cpp: ########## @@ -0,0 +1,395 @@ +/* + * 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_deduplicate_writer.h" + +#include <algorithm> +#include <optional> +#include <set> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/table/special_fields.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/scope_guard.h" +#include "paimon/core/append/append_only_writer.h" +#include "paimon/core/compact/compact_deletion_file.h" +#include "paimon/core/core_options.h" +#include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" +#include "paimon/core/index/index_file_meta.h" +#include "paimon/core/io/compact_increment.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/io/data_increment.h" +#include "paimon/core/realtime/arrow_deduplicate_realtime_store.h" +#include "paimon/core/realtime/realtime_context_impl.h" +#include "paimon/core/realtime/realtime_deduplicate_state.h" +#include "paimon/core/realtime/realtime_key_offset_index.h" +#include "paimon/core/realtime/realtime_offset_file_index_lookup.h" +#include "paimon/core/realtime/realtime_offset_utils.h" +#include "paimon/core/realtime/realtime_schema_layout.h" +#include "paimon/core/realtime/realtime_utils.h" +#include "paimon/core/utils/commit_increment.h" +#include "paimon/macros.h" +#include "paimon/realtime/realtime_context.h" + +namespace paimon { +namespace { + +template <typename T> +void AppendAll(const std::vector<std::shared_ptr<T>>& source, + std::vector<std::shared_ptr<T>>* target) { + target->insert(target->end(), source.begin(), source.end()); +} + +} // namespace + +Result<std::shared_ptr<RealtimeDeduplicateWriter>> RealtimeDeduplicateWriter::Create( + const std::map<std::string, std::string>& partition, int32_t bucket, + const std::shared_ptr<RealtimeContext>& realtime_context, + const std::shared_ptr<AppendOnlyWriter>& file_writer, + const std::shared_ptr<RealtimeSchemaLayout>& schema_layout, int64_t schema_id, + const std::vector<std::string>& deduplicate_key_fields, + const std::vector<std::shared_ptr<DataFileMeta>>& restored_data_files, + const std::shared_ptr<DataFilePathFactory>& data_file_path_factory, + const std::shared_ptr<BucketedDvMaintainer>& dv_maintainer, const CoreOptions& options, + const std::string& temp_directory, const std::shared_ptr<MemoryPool>& memory_pool) { + if (!realtime_context || !file_writer || !schema_layout || !data_file_path_factory || + !dv_maintainer || !memory_pool) { + return Status::Invalid("deduplicate real-time writer is missing a dependency"); + } + PAIMON_ASSIGN_OR_RAISE(int32_t key_position, + RealtimeUtils::GetDeduplicateKeyPosition( + schema_layout->StoreWriteSchema(), deduplicate_key_fields)); + PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(schema_layout->StoreWriteSchema())); + + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeContextImpl> context, + RealtimeContextImpl::Cast(realtime_context)); + auto write_schema = std::make_unique<ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*schema_layout->StoreWriteSchema(), write_schema.get())); + RealtimeStoreCreateRequest request{std::move(write_schema), options.ToMap(), memory_pool, + RealtimeStoreMode::DEDUPLICATE, + options.GetRealtimeStoreStatisticsMode()}; + request.deduplicate_key_fields = deduplicate_key_fields; + if (options.RealtimeSpillEnabled()) { + request.temp_directory = temp_directory; + request.file_system = options.GetFileSystem(); + } + PAIMON_ASSIGN_OR_RAISE(RealtimeStoreState store_state, + context->GetOrCreateRealtimeStore( + std::move(request), RealtimePartitionBucket(partition, bucket))); + std::shared_ptr<ArrowDeduplicateRealtimeStore> store = + std::dynamic_pointer_cast<ArrowDeduplicateRealtimeStore>(store_state.store); + if (!store || !store_state.deduplicate_state) { + return Status::Invalid("DEDUPLICATE mode requires the default Arrow store and state"); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<RealtimeOffsetFileIndexLookup> created_lookup, + RealtimeOffsetFileIndexLookup::Create( + schema_layout->CommitSchema(), schema_id, + schema_layout->StoreWriteSchema()->field(key_position), restored_data_files, + data_file_path_factory, options.GetFileSystem(), memory_pool, options)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<RealtimeOffsetFileIndexLookup> committed_lookup, + store_state.deduplicate_state->AttachFileIndexLookup(created_lookup, store)); + (void)committed_lookup; + return std::shared_ptr<RealtimeDeduplicateWriter>(new RealtimeDeduplicateWriter( + store, file_writer, schema_layout, schema_layout->StoreWriteSchema()->field(key_position), + data_file_path_factory, options.GetFileSystem(), dv_maintainer, store_state.initial_offset, + memory_pool, store_state.deduplicate_state, options.ToMap())); +} + +RealtimeDeduplicateWriter::RealtimeDeduplicateWriter( + std::shared_ptr<ArrowDeduplicateRealtimeStore> realtime_store, + std::shared_ptr<AppendOnlyWriter> file_writer, + std::shared_ptr<RealtimeSchemaLayout> schema_layout, + std::shared_ptr<arrow::Field> deduplicate_key_field, + std::shared_ptr<DataFilePathFactory> data_file_path_factory, + std::shared_ptr<FileSystem> file_system, std::shared_ptr<BucketedDvMaintainer> dv_maintainer, + int64_t next_offset, std::shared_ptr<MemoryPool> memory_pool, + std::shared_ptr<RealtimeDeduplicateState> deduplicate_state, + std::map<std::string, std::string> options) + : arrow_pool_(GetArrowPool(memory_pool)), + realtime_store_(std::move(realtime_store)), + file_writer_(std::move(file_writer)), + schema_layout_(std::move(schema_layout)), + deduplicate_key_field_(std::move(deduplicate_key_field)), + data_file_path_factory_(std::move(data_file_path_factory)), + file_system_(std::move(file_system)), + dv_maintainer_(std::move(dv_maintainer)), + memory_pool_(std::move(memory_pool)), + deduplicate_state_(std::move(deduplicate_state)), + options_(std::move(options)), + next_offset_(next_offset) {} + +Status RealtimeDeduplicateWriter::Write(std::unique_ptr<RecordBatch>&& batch) { + if (!batch || !batch->GetData()) { + return Status::Invalid("deduplicate real-time write batch is null"); + } + if (batch->GetData()->length == 0) { + return Status::OK(); + } + std::vector<RecordBatch::RowKind> row_kinds = batch->GetRowKind(); + if (row_kinds.empty()) { + row_kinds.assign(batch->GetData()->length, RecordBatch::RowKind::INSERT); + } + if (row_kinds.size() != static_cast<size_t>(batch->GetData()->length)) { + return Status::Invalid("deduplicate row-kind count does not match batch row count"); + } + + std::lock_guard<std::mutex> lock(realtime_store_mutex_); + PAIMON_ASSIGN_OR_RAISE(RealtimeOffsetUtils::ValidatedBatch validated, + RealtimeOffsetUtils::ValidateBatch( + batch.get(), schema_layout_->InputSchema(), next_offset_)); + const int32_t key_position = + schema_layout_->StoreWriteSchema()->GetFieldIndex(deduplicate_key_field_->name()); + if (key_position < 0) { + return Status::Invalid("deduplicate key field is missing from write data"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::StructArray> keys, + arrow::StructArray::Make({validated.data->field(key_position)}, {deduplicate_key_field_})); + + // Pin the complete lookup under the short store lock, then perform sidecar lookup without + // holding it. Release the building-index snapshot before mutating the index, so ordinary + // writes do not copy the entire map. The writer mutex keeps writes ordered while Advance may + // proceed. + RoaringBitmap64 previous_offsets; + { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<const KeyOffsetLookup> lookup, + realtime_store_->AcquireKeyOffsetLookup()); + PAIMON_ASSIGN_OR_RAISE(previous_offsets, lookup->LookupOffsets(keys)); + } + PAIMON_RETURN_NOT_OK(realtime_store_->WriteAndLookup(validated.data, row_kinds, + validated.offset_range, previous_offsets)); + next_offset_ = validated.offset_range.end; + has_building_data_ = true; + return Status::OK(); +} + +Status RealtimeDeduplicateWriter::SealCurrentSegment() { + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<RealtimeSegmentHandle>> segment, + realtime_store_->SealForCommit()); + std::lock_guard<std::mutex> lock(realtime_store_mutex_); + if (!segment) { + return Status::OK(); + } + if (!segment.value()) { + return Status::Invalid("deduplicate real-time store sealed a null segment"); + } + has_building_data_ = next_offset_ > segment.value()->GetOffsetRange().end; + sealed_segments_.push_back(std::move(segment.value())); + return Status::OK(); +} + +Status RealtimeDeduplicateWriter::Seal() { + std::lock_guard<std::mutex> lock(prepare_mutex_); + return SealCurrentSegment(); +} + +Result<CommitIncrement> RealtimeDeduplicateWriter::PrepareCommit(bool wait_compaction) { + std::lock_guard<std::mutex> prepare_lock(prepare_mutex_); + PAIMON_RETURN_NOT_OK(SealCurrentSegment()); + std::vector<std::shared_ptr<RealtimeSegmentHandle>> segments; + { + std::lock_guard<std::mutex> store_lock(realtime_store_mutex_); + segments.swap(sealed_segments_); + } + + std::vector<std::shared_ptr<DataFileMeta>> new_files; + std::vector<std::shared_ptr<IndexFileMeta>> new_index_files; + std::vector<std::shared_ptr<IndexFileMeta>> deleted_index_files; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<RealtimeOffsetFileIndexLookup> committed_file_lookup, + deduplicate_state_->AcquireCommittedFileLookup()); + std::vector<std::shared_ptr<DataFileMeta>> lookup_data_files = + committed_file_lookup->DataFiles(); + std::set<std::string> committed_file_names; + for (const std::shared_ptr<DataFileMeta>& data_file : lookup_data_files) { + committed_file_names.insert(data_file->file_name); + } + prepared_data_files_.erase( + std::remove_if(prepared_data_files_.begin(), prepared_data_files_.end(), + [&committed_file_names](const std::shared_ptr<DataFileMeta>& data_file) { + return committed_file_names.count(data_file->file_name) > 0; + }), + prepared_data_files_.end()); + lookup_data_files.insert(lookup_data_files.end(), prepared_data_files_.begin(), + prepared_data_files_.end()); Review Comment: [P1] Writer handoff loses prepared files needed for deletion vectors `prepared_data_files_` belongs to this writer, while the shared context retains sealed key/offset state across writer handoff. On a new table, W1 can write `k@0`, prepare F1, and close before commit; W2 can reuse the context, write `k@1`, and prepare F2. W2 finds offset 0 in the retained key lookup, but the file list here contains neither F1 nor any committed files, so no DV is generated for F1. Committing the accumulated results then leaves both versions visible on disk. This lifecycle is explicitly exercised for ordinary append by `TestCloseWriterAllowsContextReuseByLaterWriter`; the dedup handoff test commits before closing and does not cover this window. -- 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]
