From 9b0b5cece6638886f6ef9489ec8d09f222ea0a95 Mon Sep 17 00:00:00 2001 From: JeffZhou <17023790+HaHaJeff@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:31:32 +0800 Subject: [PATCH] feat(realtime): add primary-key realtime store --- src/paimon/CMakeLists.txt | 2 + src/paimon/common/utils/arrow/arrow_utils.cpp | 16 + src/paimon/common/utils/arrow/arrow_utils.h | 3 + src/paimon/common/utils/arrow/mem_utils.cpp | 41 ++ src/paimon/common/utils/arrow/mem_utils.h | 6 + .../core/realtime/arrow_realtime_store.cpp | 25 +- .../realtime/arrow_realtime_store_factory.cpp | 7 +- .../realtime/primary_key_realtime_store.cpp | 309 ++++++++++++ .../realtime/primary_key_realtime_store.h | 61 +++ .../primary_key_realtime_store_test.cpp | 451 ++++++++++++++++++ 10 files changed, 899 insertions(+), 22 deletions(-) create mode 100644 src/paimon/core/realtime/primary_key_realtime_store.cpp create mode 100644 src/paimon/core/realtime/primary_key_realtime_store.h create mode 100644 src/paimon/core/realtime/primary_key_realtime_store_test.cpp diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 1749634f4..b8e5ec3e8 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -385,6 +385,7 @@ set(PAIMON_CORE_SRCS core/operation/write_restore.cpp core/realtime/arrow_realtime_store.cpp core/realtime/arrow_realtime_store_factory.cpp + core/realtime/primary_key_realtime_store.cpp core/realtime/realtime_append_only_writer.cpp core/realtime/realtime_context.cpp core/realtime/realtime_context_impl.cpp @@ -791,6 +792,7 @@ if(PAIMON_BUILD_TESTS) core/manifest/index_manifest_file_handler_test.cpp core/memory/writer_memory_manager_test.cpp core/realtime/arrow_realtime_store_test.cpp + core/realtime/primary_key_realtime_store_test.cpp core/realtime/realtime_context_test.cpp core/realtime/realtime_reader_test.cpp core/mergetree/levels_test.cpp diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index 97bb77813..1ba4c153b 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -364,6 +364,22 @@ void ArrowUtils::TraverseArray(const std::shared_ptr& array) { } } +uint64_t ArrowUtils::GetArrayMemoryUsage(const std::shared_ptr& data) { + uint64_t result = 0; + for (const std::shared_ptr& buffer : data->buffers) { + if (buffer) { + result += static_cast(buffer->size()); + } + } + for (const std::shared_ptr& child : data->child_data) { + result += GetArrayMemoryUsage(child); + } + if (data->dictionary) { + result += GetArrayMemoryUsage(data->dictionary); + } + return result; +} + bool ArrowUtils::EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type) { if (type->id() != other_type->id() || type->num_fields() != other_type->num_fields()) { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index d82d84d88..59db09815 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include "arrow/api.h" @@ -48,6 +49,8 @@ class PAIMON_EXPORT ArrowUtils { // avoid subsequent multi-threading problems. static void TraverseArray(const std::shared_ptr& array); + static uint64_t GetArrayMemoryUsage(const std::shared_ptr& data); + static Result> RemoveFieldFromStructArray( const std::shared_ptr& struct_array, const std::string& field_name); diff --git a/src/paimon/common/utils/arrow/mem_utils.cpp b/src/paimon/common/utils/arrow/mem_utils.cpp index 7e8986be5..1e9695332 100644 --- a/src/paimon/common/utils/arrow/mem_utils.cpp +++ b/src/paimon/common/utils/arrow/mem_utils.cpp @@ -24,12 +24,31 @@ #include #include +#include "arrow/c/abi.h" +#include "arrow/c/helpers.h" #include "arrow/memory_pool.h" #include "arrow/status.h" #include "fmt/format.h" #include "paimon/memory/memory_pool.h" namespace paimon { +namespace { + +struct ArrowArrayPrivateData { + void (*release)(ArrowArray*); + void* private_data; + std::shared_ptr arrow_pool; +}; + +void ReleaseArrowArray(ArrowArray* array) { + std::unique_ptr data( + static_cast(array->private_data)); + array->release = data->release; + array->private_data = data->private_data; + array->release(array); +} + +} // namespace class ArrowMemPoolAdaptor : public arrow::MemoryPool { public: @@ -107,4 +126,26 @@ std::unique_ptr GetArrowPool(const std::shared_ptr(pool); } +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool) { + if (!array || !array->release) { + return Status::Invalid("cannot retain Arrow array memory pool"); + } + if (!arrow_pool) { + ArrowArrayRelease(array); + return Status::Invalid("cannot retain Arrow array memory pool"); + } + std::unique_ptr data; + try { + data = std::make_unique( + ArrowArrayPrivateData{array->release, array->private_data, arrow_pool}); + } catch (const std::bad_alloc&) { + ArrowArrayRelease(array); + return Status::OutOfMemory("failed to retain Arrow array memory pool"); + } + array->private_data = data.release(); + array->release = ReleaseArrowArray; + return Status::OK(); +} + } // namespace paimon diff --git a/src/paimon/common/utils/arrow/mem_utils.h b/src/paimon/common/utils/arrow/mem_utils.h index 96b59e3e8..214bb4509 100644 --- a/src/paimon/common/utils/arrow/mem_utils.h +++ b/src/paimon/common/utils/arrow/mem_utils.h @@ -23,11 +23,17 @@ #include "arrow/memory_pool.h" #include "paimon/memory/memory_pool.h" +#include "paimon/status.h" #include "paimon/visibility.h" +struct ArrowArray; + namespace paimon { PAIMON_EXPORT std::unique_ptr GetArrowPool( const std::shared_ptr& pool); +Status RetainArrowArrayMemoryPool(ArrowArray* array, + const std::shared_ptr& arrow_pool); + } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index 18087a29f..1136243e8 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -33,6 +33,7 @@ #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/row_kind.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/projected_array.h" @@ -43,22 +44,6 @@ namespace paimon { namespace { -uint64_t GetArrayMemoryUsage(const std::shared_ptr& data) { - uint64_t result = 0; - for (const std::shared_ptr& buffer : data->buffers) { - if (buffer) { - result += static_cast(buffer->size()); - } - } - for (const std::shared_ptr& child : data->child_data) { - result += GetArrayMemoryUsage(child); - } - if (data->dictionary) { - result += GetArrayMemoryUsage(data->dictionary); - } - return result; -} - bool SupportsMinMax(const std::shared_ptr& type) { switch (type->id()) { case arrow::Type::BOOL: @@ -393,11 +378,11 @@ Status ArrowRealtimeStore::Write(RealtimeWriteBatch&& write_batch) { if (building_range_ && write_batch.offset_range.begin != building_range_->end) { return Status::Invalid("real-time offset ranges must be contiguous"); } - uint64_t memory_usage = GetArrayMemoryUsage(struct_array->data()); + uint64_t memory_usage = ArrowUtils::GetArrayMemoryUsage(struct_array->data()); if (statistics) { - memory_usage += GetArrayMemoryUsage(statistics->min_values->data()) + - GetArrayMemoryUsage(statistics->max_values->data()) + - GetArrayMemoryUsage(statistics->null_counts->data()); + memory_usage += ArrowUtils::GetArrayMemoryUsage(statistics->min_values->data()) + + ArrowUtils::GetArrayMemoryUsage(statistics->max_values->data()) + + ArrowUtils::GetArrayMemoryUsage(statistics->null_counts->data()); } building_memory_usage_ += memory_usage; building_batches_.push_back( diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index edd8cfeee..dff12b589 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -25,6 +25,7 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/primary_key_realtime_store.h" #include "paimon/macros.h" namespace paimon { @@ -48,8 +49,10 @@ Result> ArrowRealtimeStoreFactory::Create( request.memory_pool, arrow_pool); } case RealtimeStoreMode::PRIMARY_KEY: { - return Status::NotImplemented( - "primary-key real-time store support is not installed"); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(imported_schema, request.memory_pool)); + return std::shared_ptr(std::move(store)); } } return Status::Invalid("invalid real-time store mode: ", static_cast(request.mode)); diff --git a/src/paimon/core/realtime/primary_key_realtime_store.cpp b/src/paimon/core/realtime/primary_key_realtime_store.cpp new file mode 100644 index 000000000..421eb0c8d --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.cpp @@ -0,0 +1,309 @@ +/* + * 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 +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.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/scope_guard.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" + +namespace paimon { + +namespace { + +struct StoredBatch { + std::shared_ptr data; + OffsetRange offset_range; + uint64_t memory_usage; +}; + +class Segment final : public RealtimeSegmentHandle { + public: + Segment(const OffsetRange& range, std::vector&& batches) + : range_(range), batches_(std::move(batches)) {} + + OffsetRange GetOffsetRange() const override { + return range_; + } + const std::vector& Batches() const { + return batches_; + } + + private: + OffsetRange range_; + std::vector batches_; +}; + +class ReadView final : public RealtimeReadView { + public: + explicit ReadView(std::vector>&& segments) + : segments_(std::move(segments)) { + if (!segments_.empty()) { + range_ = OffsetRange(segments_.front()->GetOffsetRange().begin, + segments_.back()->GetOffsetRange().end); + } + } + + std::optional GetOffsetRange() const override { + return range_; + } + const std::vector>& Segments() const { + return segments_; + } + + private: + std::vector> segments_; + std::optional range_; +}; + +class StoredBatchReader final : public BatchReader { + public: + explicit StoredBatchReader(const StoredBatch& batch, + std::shared_ptr arrow_pool) + : arrow_pool_(std::move(arrow_pool)), + data_(batch.data), + metrics_(std::make_shared()) {} + + Result NextBatch() override { + if (!data_) { + return MakeEofBatch(); + } + auto array = std::make_unique(); + auto schema = std::make_unique(); + ScopeGuard export_guard([array_ptr = array.get(), schema_ptr = schema.get()]() { + ArrowArrayRelease(array_ptr); + ArrowSchemaRelease(schema_ptr); + }); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr record_batch, + arrow::RecordBatch::FromStructArray(data_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr normalized_batch, + ArrowUtils::NormalizeRecordBatchOffsets(record_batch, arrow_pool_.get())); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportRecordBatch(*normalized_batch, array.get(), schema.get())); + PAIMON_RETURN_NOT_OK(RetainArrowArrayMemoryPool(array.get(), arrow_pool_)); + data_.reset(); + arrow_pool_.reset(); + export_guard.Release(); + return ReadBatch(std::move(array), std::move(schema)); + } + + std::shared_ptr GetReaderMetrics() const override { + return metrics_; + } + void Close() override { + data_.reset(); + arrow_pool_.reset(); + } + + private: + std::shared_ptr arrow_pool_; + std::shared_ptr data_; + std::shared_ptr metrics_; +}; + +} // namespace + +class PrimaryKeyRealtimeStore::Impl { + public: + Impl(std::shared_ptr transport_schema, + std::shared_ptr arrow_pool) + : transport_schema_(std::move(transport_schema)), arrow_pool_(std::move(arrow_pool)) {} + + 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 (write_batch.offset_range.begin < 0 || write_batch.offset_range.Count() != row_count || + row_count <= 0) { + return Status::Invalid("PK real-time offset range does not match batch row count"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(write_batch.batch->GetData(), + arrow::struct_(transport_schema_->fields()))); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("PK real-time transport batch is not a StructArray"); + } + std::shared_ptr transport = + checked_pointer_cast(array); + std::lock_guard lock(mutex_); + building_.push_back(StoredBatch{transport, write_batch.offset_range, + ArrowUtils::GetArrayMemoryUsage(transport->data())}); + building_memory_usage_ += building_.back().memory_usage; + return Status::OK(); + } + + Result>> SealForCommit() { + std::lock_guard lock(mutex_); + if (building_.empty()) { + return std::optional>(); + } + OffsetRange range(building_.front().offset_range.begin, building_.back().offset_range.end); + std::shared_ptr segment = std::make_shared(range, std::move(building_)); + sealed_.push_back(segment); + building_.clear(); + building_memory_usage_ = 0; + return std::optional>(std::move(segment)); + } + + Result>> CreateCommitReaders( + const std::shared_ptr& handle) { + std::shared_ptr segment = std::dynamic_pointer_cast(handle); + if (!segment) { + return Status::Invalid("segment was not created by the PK real-time store"); + } + std::vector> readers; + readers.reserve(segment->Batches().size()); + for (const StoredBatch& batch : segment->Batches()) { + readers.push_back(std::make_unique(batch, arrow_pool_)); + } + return readers; + } + + Result> AcquireReadView() { + std::lock_guard lock(mutex_); + std::vector> segments = sealed_; + if (!building_.empty()) { + OffsetRange range(building_.front().offset_range.begin, + building_.back().offset_range.end); + segments.push_back( + std::make_shared(range, std::vector(building_))); + } + return std::make_shared(std::move(segments)); + } + + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t, + const RealtimeQueryContext& context) { + std::shared_ptr typed = std::dynamic_pointer_cast(view); + if (!typed) { + return Status::Invalid("read view was not created by the PK real-time store"); + } + if (context.read_schema == nullptr || context.read_schema->release == nullptr) { + return Status::Invalid("PK real-time query read schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, + arrow::ImportSchema(context.read_schema)); + std::vector> readers; + for (const std::shared_ptr& segment : typed->Segments()) { + for (const StoredBatch& batch : segment->Batches()) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr projected, + NestedProjectionUtils::AlignArrayToReadType( + batch.data, arrow::struct_(read_schema->fields()), arrow_pool_.get())); + if (!projected || projected->type_id() != arrow::Type::STRUCT) { + return Status::Invalid( + "PK memory query projection did not produce a StructArray"); + } + StoredBatch query_batch{checked_pointer_cast(projected), + batch.offset_range, /*memory_usage=*/0}; + readers.push_back(std::make_unique(query_batch, arrow_pool_)); + } + } + return readers; + } + + Status AdvanceCommittedOffset(int64_t committed_end_offset) { + std::lock_guard lock(mutex_); + auto first_retained = std::find_if( + sealed_.begin(), sealed_.end(), [committed_end_offset](const auto& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + sealed_.erase(sealed_.begin(), first_retained); + return Status::OK(); + } + + uint64_t GetMemoryUsage() const { + std::lock_guard lock(mutex_); + uint64_t total = building_memory_usage_; + for (const std::shared_ptr& segment : sealed_) { + for (const StoredBatch& batch : segment->Batches()) { + total += batch.memory_usage; + } + } + return total; + } + + private: + std::shared_ptr transport_schema_; + std::shared_ptr arrow_pool_; + mutable std::mutex mutex_; + std::vector building_; + std::vector> sealed_; + uint64_t building_memory_usage_ = 0; +}; + +PrimaryKeyRealtimeStore::PrimaryKeyRealtimeStore(std::unique_ptr&& impl) + : impl_(std::move(impl)) {} +PrimaryKeyRealtimeStore::~PrimaryKeyRealtimeStore() = default; + +Result> PrimaryKeyRealtimeStore::Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& memory_pool) { + if (!memory_pool) { + return Status::Invalid("PK real-time store memory pool is null"); + } + std::shared_ptr arrow_pool = GetArrowPool(memory_pool); + return std::shared_ptr(new PrimaryKeyRealtimeStore( + std::make_unique(transport_schema, std::move(arrow_pool)))); +} +Status PrimaryKeyRealtimeStore::Write(RealtimeWriteBatch&& batch) { + return impl_->Write(std::move(batch)); +} +Result>> +PrimaryKeyRealtimeStore::SealForCommit() { + return impl_->SealForCommit(); +} +Result>> PrimaryKeyRealtimeStore::CreateCommitReaders( + const std::shared_ptr& segment) { + return impl_->CreateCommitReaders(segment); +} +Result> PrimaryKeyRealtimeStore::AcquireReadView() { + return impl_->AcquireReadView(); +} +Result>> PrimaryKeyRealtimeStore::CreateQueryReaders( + const std::shared_ptr& view, int64_t offset, + const RealtimeQueryContext& context) { + return impl_->CreateQueryReaders(view, offset, context); +} +Status PrimaryKeyRealtimeStore::AdvanceCommittedOffset(int64_t committed_end_offset) { + return impl_->AdvanceCommittedOffset(committed_end_offset); +} +uint64_t PrimaryKeyRealtimeStore::GetMemoryUsage() const { + return impl_->GetMemoryUsage(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store.h b/src/paimon/core/realtime/primary_key_realtime_store.h new file mode 100644 index 000000000..46c0fe8f7 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store.h @@ -0,0 +1,61 @@ +/* + * 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 "paimon/realtime/realtime_store.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { + +class MemoryPool; + +/// Internal in-memory implementation of the default primary-key `RealtimeStore`. +class PrimaryKeyRealtimeStore final : public RealtimeStore { + public: + static Result> Create( + const std::shared_ptr& transport_schema, + const std::shared_ptr& memory_pool); + + ~PrimaryKeyRealtimeStore() override; + + Status Write(RealtimeWriteBatch&& batch) override; + Result>> SealForCommit() override; + Result>> CreateCommitReaders( + const std::shared_ptr& segment) override; + Result> AcquireReadView() override; + Result>> CreateQueryReaders( + const std::shared_ptr& view, int64_t offset_begin, + const RealtimeQueryContext& context) override; + Status AdvanceCommittedOffset(int64_t committed_end_offset) override; + uint64_t GetMemoryUsage() const override; + + private: + class Impl; + explicit PrimaryKeyRealtimeStore(std::unique_ptr&& impl); + + std::unique_ptr impl_; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp new file mode 100644 index 000000000..7d6fa72f5 --- /dev/null +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -0,0 +1,451 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/realtime/arrow_realtime_store_factory.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::shared_ptr FieldWithId(const std::string& name, + const std::shared_ptr& type, + int32_t field_id, bool nullable = true) { + return DataField::ConvertDataFieldToArrowField( + DataField(field_id, arrow::field(name, type, nullable))) + ->WithNullable(nullable); +} + +std::shared_ptr TransportSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField( + DataField(1, arrow::field("value", arrow::utf8())))}); +} + +std::shared_ptr NestedTransportSchema() { + return arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + DataField::ConvertDataFieldToArrowField(DataField(0, arrow::field("id", arrow::int64()))), + DataField::ConvertDataFieldToArrowField(DataField( + 1, + arrow::field("value", + arrow::struct_({arrow::field("name", arrow::utf8()), + arrow::field("items", arrow::list(arrow::int32()))}))))}); +} + +std::unique_ptr MakeBatch(const std::string& json) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(TransportSchema()->fields()), json) + .ValueOrDie(); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +std::unique_ptr MakeSlicedBatch(const std::shared_ptr& schema, + const std::string& json, int64_t offset, + int64_t length) { + std::shared_ptr array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie() + ->Slice(offset, length); + auto c_array = std::make_unique(); + EXPECT_TRUE(arrow::ExportArray(*array, c_array.get()).ok()); + return RecordBatchBuilder(c_array.get()).Finish().value(); +} + +void AssertOffsetsZero(const ArrowArray* array) { + ASSERT_NE(nullptr, array); + ASSERT_EQ(0, array->offset); + for (int64_t child = 0; child < array->n_children; ++child) { + AssertOffsetsZero(array->children[child]); + } + if (array->dictionary) { + AssertOffsetsZero(array->dictionary); + } +} + +Result ReadJson(const std::vector>& readers) { + std::vector> batches; + for (const std::unique_ptr& reader : readers) { + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + batches.push_back(std::move(array)); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr result, + arrow::Concatenate(batches)); + return result->ToString(); +} + +class TestingMemoryPool final : public MemoryPool { + public: + void* Malloc(uint64_t size, uint64_t alignment) override { + return delegate_->Malloc(size, alignment); + } + + void* Realloc(void* pointer, size_t old_size, size_t new_size, uint64_t alignment) override { + return delegate_->Realloc(pointer, old_size, new_size, alignment); + } + + void Free(void* pointer, uint64_t size) override { + delegate_->Free(pointer, size); + } + + void Free(void* pointer, uint64_t size, uint64_t alignment) override { + delegate_->Free(pointer, size, alignment); + } + + uint64_t CurrentUsage() const override { + return delegate_->CurrentUsage(); + } + + uint64_t MaxMemoryUsage() const override { + return delegate_->MaxMemoryUsage(); + } + + private: + std::unique_ptr delegate_ = GetMemoryPool(); +}; + +TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_FALSE(segment.has_value()); + ASSERT_NOK_WITH_MSG(store->Write(RealtimeWriteBatch{nullptr, OffsetRange(0, 0)}), + "write batch is null"); + ASSERT_NOK_WITH_MSG( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 1, "one"]])"), OffsetRange(0, 0)}), + "offset range does not match batch row count"); + + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[0, 1, 0, 1, "one"], [0, 2, 1, 2, "two"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 3, 2, 3, "three"]])"), OffsetRange(2, 3)})); + + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_EQ(OffsetRange(0, 3), segment.value()->GetOffsetRange()); + ASSERT_GT(store->GetMemoryUsage(), 0); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 4, 3, 4, "four"]])"), OffsetRange(3, 4)})); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[1, 6, 1, 1, "before"], [0, 5, 0, 3, "three"]])"), OffsetRange(0, 2)})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[2, 7, 2, 2, "after"]])"), OffsetRange(2, 3)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_EQ( + "-- is_valid: all not null\n-- child 0 type: int8\n [\n 1,\n 0,\n 2\n ]\n-- " + "child 1 type: int64\n [\n 6,\n 5,\n 7\n ]\n-- child 2 type: int64\n [\n " + "1,\n 0,\n 2\n ]\n-- child 3 type: int64\n [\n 1,\n 3,\n 2\n ]\n-- child " + "4 type: string\n [\n \"before\",\n \"three\",\n \"after\"\n ]", + actual); + for (const std::unique_ptr& reader : readers) { + reader->Close(); + } +} + +void AssertSlicedBatch(BatchReader* reader) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + ASSERT_EQ(2, batch.first->length); + AssertOffsetsZero(batch.first.get()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + std::shared_ptr values = checked_pointer_cast(array); + ASSERT_EQ(2, checked_pointer_cast(values->field(3))->Value(0)); + ASSERT_EQ(3, checked_pointer_cast(values->field(3))->Value(1)); + std::shared_ptr nested = + checked_pointer_cast(values->field(4)); + ASSERT_EQ("two", checked_pointer_cast(nested->field(0))->GetString(0)); + ASSERT_EQ("three", checked_pointer_cast(nested->field(0))->GetString(1)); + std::shared_ptr items = + checked_pointer_cast(nested->field(1)); + std::shared_ptr first_items = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(3, first_items->Value(0)); + ASSERT_EQ(4, first_items->Value(1)); + std::shared_ptr second_items = + checked_pointer_cast(items->value_slice(1)); + ASSERT_EQ(5, second_items->Value(0)); + ASSERT_EQ(6, second_items->Value(1)); + ASSERT_OK_AND_ASSIGN(batch, reader->NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(batch)); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { + std::shared_ptr schema = NestedTransportSchema(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + schema, + R"([[0, 1, 0, 1, ["one", [1, 2]]], [0, 2, 1, 2, ["two", [3, 4]]], [0, 3, 2, 3, ["three", [5, 6]]], [0, 4, 3, 4, ["four", [7, 8]]]])", + 1, 2), + OffsetRange(0, 2)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateCommitReaders(segment.value())); + ASSERT_EQ(1, readers.size()); + AssertSlicedBatch(readers[0].get()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(readers, store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + AssertSlicedBatch(readers[0].get()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 5, 2, "two"]])"), OffsetRange(5, 6)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 2, 6, 3, "three"]])"), OffsetRange(6, 7)})); + ASSERT_OK_AND_ASSIGN(segment, store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr retained_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), retained_view->GetOffsetRange()); + + const uint64_t initial_memory_usage = store->GetMemoryUsage(); + ASSERT_GT(initial_memory_usage, 0); + ASSERT_OK(store->AdvanceCommittedOffset(4)); + ASSERT_EQ(initial_memory_usage, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(4, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(5)); + ASSERT_LT(store->GetMemoryUsage(), initial_memory_usage); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(5, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(6)); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_EQ(OffsetRange(6, 7), current_view->GetOffsetRange()); + ASSERT_OK(store->AdvanceCommittedOffset(7)); + ASSERT_EQ(0, store->GetMemoryUsage()); + ASSERT_OK_AND_ASSIGN(current_view, store->AcquireReadView()); + ASSERT_FALSE(current_view->GetOffsetRange().has_value()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(retained_view, /*offset_begin=*/0, context)); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); + ASSERT_NE(std::string::npos, actual.find("\"three\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(TransportSchema(), GetDefaultPool())); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 2, "two"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 2, 1, 1, "one"]])"), OffsetRange(1, 2)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*TransportSchema(), c_schema.get()).ok()); + RealtimeQueryContext context{/*read_schema=*/c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(2, readers.size()); + ASSERT_OK_AND_ASSIGN(std::string actual, ReadJson(readers)); + ASSERT_NE(std::string::npos, actual.find("\"one\"")); + ASSERT_NE(std::string::npos, actual.find("\"two\"")); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryPoolOutlivesStoreReaderAndExport) { + const std::shared_ptr stored_schema = TransportSchema(); + std::shared_ptr pool = std::make_shared(); + std::weak_ptr pool_lifetime = pool; + auto write_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, write_schema.get()).ok()); + ArrowRealtimeStoreFactory factory; + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + factory.Create(RealtimeStoreCreateRequest{ + std::move(write_schema), + /*options=*/{}, pool, RealtimeStoreMode::PRIMARY_KEY})); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[0, 1, 0, 7, "seven"]])"), OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*stored_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_EQ(1, readers.size()); + view.reset(); + store.reset(); + pool.reset(); + ASSERT_FALSE(pool_lifetime.expired()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr imported = std::move(import_result).ValueOrDie(); + readers.clear(); + ASSERT_FALSE(pool_lifetime.expired()); + imported.reset(); + ASSERT_TRUE(pool_lifetime.expired()); +} + +TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { + const std::shared_ptr stored_profile_a = + FieldWithId("profile_a", arrow::int32(), 30); + const std::shared_ptr stored_a = FieldWithId("a", arrow::int32(), 10); + const std::shared_ptr stored_b = FieldWithId("b", arrow::int32(), 11); + const std::shared_ptr stored_x = FieldWithId("x", arrow::int32(), 20); + const std::shared_ptr stored_y = FieldWithId("y", arrow::int32(), 21); + arrow::FieldVector stored_fields = { + DataField::ConvertDataFieldToArrowField(SpecialFields::ValueKind())->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::SequenceNumber()) + ->WithNullable(false), + DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + FieldWithId("id", arrow::int64(), 0), + FieldWithId("profile", arrow::struct_({stored_profile_a}), 1), + FieldWithId("items", arrow::list(arrow::struct_({stored_a, stored_b})), 2), + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_x, stored_y})), 3)}; + std::shared_ptr stored_schema = arrow::schema(std::move(stored_fields)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + PrimaryKeyRealtimeStore::Create(stored_schema, GetDefaultPool())); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeSlicedBatch( + stored_schema, + R"([[0, 1, 0, 6, [5], [[1, 2]], [["before", [3, 4]]]], [0, 2, 1, 7, [50], [[100, 200], null], [["k1", [7, 8]], ["k2", null]]], [0, 3, 2, 8, [500], [[9, 10]], [["after", [11, 12]]]]])", + 1, 1), + OffsetRange(0, 1)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + + arrow::FieldVector requested_fields(stored_schema->fields().begin(), + stored_schema->fields().begin() + 3); + requested_fields.push_back(FieldWithId("profile", arrow::struct_({stored_profile_a}), 1)); + requested_fields.push_back( + FieldWithId("items", arrow::list(arrow::struct_({stored_b, stored_a})), 2)); + requested_fields.push_back( + FieldWithId("attrs", arrow::map(arrow::utf8(), arrow::struct_({stored_y, stored_x})), 3)); + std::shared_ptr requested_schema = arrow::schema(std::move(requested_fields)); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*requested_schema, c_schema.get()).ok()); + RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + /*enable_predicate_pushdown=*/false}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, /*offset_begin=*/0, context)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); + arrow::Result> import_result = + arrow::ImportArray(batch.first.get(), batch.second.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + std::shared_ptr array = std::move(import_result).ValueOrDie(); + ASSERT_TRUE(array->type()->Equals(arrow::struct_(requested_schema->fields()))); + std::shared_ptr projected = checked_pointer_cast(array); + const std::shared_ptr profile = + checked_pointer_cast(projected->field(3)); + ASSERT_EQ(50, checked_pointer_cast(profile->field(0))->Value(0)); + const std::shared_ptr items = + checked_pointer_cast(projected->field(4)); + const std::shared_ptr item_values = + checked_pointer_cast(items->value_slice(0)); + ASSERT_EQ(200, checked_pointer_cast(item_values->field(0))->Value(0)); + ASSERT_EQ(100, checked_pointer_cast(item_values->field(1))->Value(0)); + ASSERT_TRUE(item_values->IsNull(1)); + + const std::shared_ptr attrs = + checked_pointer_cast(projected->field(5)); + const int64_t attr_offset = attrs->value_offset(0); + const int64_t attr_length = attrs->value_length(0); + const std::shared_ptr attr_keys = + checked_pointer_cast(attrs->keys()->Slice(attr_offset, attr_length)); + ASSERT_EQ("k1", attr_keys->GetString(0)); + const std::shared_ptr attr_values = + checked_pointer_cast(attrs->items()->Slice(attr_offset, attr_length)); + ASSERT_EQ(8, checked_pointer_cast(attr_values->field(0))->Value(0)); + ASSERT_EQ(7, checked_pointer_cast(attr_values->field(1))->Value(0)); + ASSERT_TRUE(attr_values->IsNull(1)); +} + +} // namespace +} // namespace paimon::test