diff --git a/docs/source/user_guide/realtime.rst b/docs/source/user_guide/realtime.rst index 5033d547a..f291605d0 100644 --- a/docs/source/user_guide/realtime.rst +++ b/docs/source/user_guide/realtime.rst @@ -115,6 +115,11 @@ Common Options - ``none`` - ``none`` or ``full``. Full mode collects per-batch min/max and null counts for conservative predicate pruning, including retained spill metadata. + * - ``realtime.index..columns`` + - Unset + - Comma-separated fields indexed in memory by the default Arrow store. + Available in append and append deduplicate modes; independent of + ``realtime.store.stats-mode``. * - ``realtime.deduplicate-key-fields`` - Unset - Selects append deduplicate mode. Although expressed as a comma-separated @@ -398,6 +403,37 @@ context's reclaimed progress for relevant active stores. It does not restore an old in-memory state and is not unrestricted historical time travel. For a disk-only historical query, do not attach the real-time context. +In-memory Predicate Indexes +-------------------------- + +The default Arrow store supports optional predicate indexes on key or value fields +in append and append deduplicate modes. PK modes are unsupported. For example: + +.. code-block:: text + + realtime.enabled = true + realtime.index.hash.columns = id,payload + realtime.store.stats-mode = full + +The built-in ``hash`` plugin is **experimental**; its implementation and supported +capabilities may change. It supports equality and ``IN`` predicates on boolean, +integer, floating-point, string, binary, and date fields. Decimal, timestamp, and +nested types are currently unsupported. Use a comma-separated list to index +multiple fields; indexing ``_REALTIME_OFFSET`` is unsupported. + +Note that this index does not currently handle hash collisions. + +Indexes and batch statistics work independently and can be enabled together. +For ``id = 7 AND score > 100``, an index supporting equality on ``id`` selects +candidate offsets; statistics can then skip batches whose maximum ``score`` +is at most 100. +Statistics cover entire batches, so mixed batches still need row-level +filtering. Keep predicate fields in the read schema and enable +``ReadContextBuilder::EnablePredicateFilter(true)`` for exact filtering. + +In-memory and disk indexes are configured independently. Use +``file-index..columns`` to enable indexes for committed files. + Implementation by Mode ---------------------- @@ -406,9 +442,10 @@ Append ``RealtimeAppendOnlyWriter`` validates offsets and writes to an ``APPEND_ONLY`` store. The default ``ArrowRealtimeStore`` retains building -batches and immutable sealed segments in write order. Query readers project -the requested fields, optionally prune stored batches using statistics, and -apply the visible offset range before removing the hidden offset column. +batches and immutable sealed segments in write order. Query readers use +optional predicate indexes and batch statistics to prune candidates, project +the requested fields, and apply the visible offset range before removing the +hidden offset column. Preparing reads sealed segments through the ordinary ``AppendOnlyWriter``, removing the internal offset prefix before writing table files. Queries do not @@ -467,8 +504,10 @@ are marked deleted, and within-batch duplicates are resolved in row order. Add rows are appended to the store; delete events update only the key/DV state, without persisting a delete row. -Key indexes and DV bitmaps use copy-on-write snapshots, so later writes do not -change the state held by existing lookups, read views, or seal checkpoints. +Deduplicate key-to-latest-offset lookups and DV bitmaps use copy-on-write +snapshots, so later writes do not change the state held by existing lookups, +read views, or seal checkpoints. The optional predicate indexes share +append-only postings instead. Memory reads use the read view's DV to exclude rows whose offsets are marked deleted. @@ -511,8 +550,9 @@ still terminal for the owning writer/context. Old views can retain the memory representation after spill or reclamation; new views use the published representation. Spill files are deleted when the store and all remaining handles/readers have released them. They are temporary -storage, not a recoverable log after process restart. Deduplicate key indexes -and DVs remain in memory; full-mode statistics also remain resident after spill. +storage, not a recoverable log after process restart. Predicate indexes, +deduplicate key indexes, DVs, and full-mode statistics remain in memory after +spill. Predicate index memory is included in store memory accounting. ``RealtimeContext::GetMetrics()`` and real-time writer metrics expose: @@ -596,8 +636,8 @@ Current restrictions include: * Append deduplicate requires a single non-null scalar key, a persisted offset column, exact offset file-index reads, and non-bitmap64 DVs. Existing ordinary append files are not migrated automatically. -* Spill covers Arrow data, not key indexes or DVs, and provides neither a hard - memory limit nor crash recovery. +* Spill covers Arrow data, not predicate indexes, key indexes, or DVs, and + provides neither a hard memory limit nor crash recovery. Code Map -------- diff --git a/include/paimon/api.h b/include/paimon/api.h index a237c7691..e339aaeb5 100644 --- a/include/paimon/api.h +++ b/include/paimon/api.h @@ -42,6 +42,9 @@ #include "paimon/write_context.h" // IWYU pragma: export // IWYU pragma: begin_exports +#include "paimon/realtime/index/realtime_index.h" +#include "paimon/realtime/index/realtime_index_factory.h" +#include "paimon/realtime/index/realtime_index_reader.h" #include "paimon/realtime/realtime_context.h" #include "paimon/realtime/realtime_snapshot_offsets.h" #include "paimon/realtime/realtime_store.h" diff --git a/include/paimon/file_store_write.h b/include/paimon/file_store_write.h index 5e6b253c5..ae01dd080 100644 --- a/include/paimon/file_store_write.h +++ b/include/paimon/file_store_write.h @@ -41,6 +41,9 @@ class RecordBatch; class WriteContext; /// Interface for write operations in a file store. +/// +/// A `RealtimeContext` supports only one active writer for each partition-bucket. Close the +/// previous writer before creating a replacement that accesses the same partition-bucket. class PAIMON_EXPORT FileStoreWrite { public: /// Create an instance of `FileStoreWrite`. @@ -134,6 +137,10 @@ class PAIMON_EXPORT FileStoreWrite { /// /// The writer loads the snapshot's partition-bucket offsets and releases sealed memory that is /// fully covered by disk. Calling this method on a non-real-time writer returns an error. + /// After a successful refresh, fully covered segments are also removed from this writer's + /// pending prepare queue. Existing read views continue to pin the resources they reference. + /// Snapshot installation and queue cleanup are serialized with `Seal()` and + /// `PrepareCommitWithProgress()`. /// /// With `WriteContextBuilder::WithCatalog()`, the latest snapshot can be loaded from the /// catalog. Older requested snapshots must be available on the file system. @@ -162,8 +169,9 @@ class PAIMON_EXPORT FileStoreWrite { /// Releases resources owned by this writer. /// - /// Closing a real-time writer with data not covered by a successful - /// `PrepareCommitWithProgress()` returns an error and invalidates its `RealtimeContext`. + /// Closing a real-time writer with data covered by neither a successful + /// `PrepareCommitWithProgress()` nor a refreshed committed snapshot returns an error and + /// invalidates its `RealtimeContext`. /// The caller must rebuild both objects and recover input from the durable snapshot offset. virtual Status Close() = 0; }; diff --git a/include/paimon/realtime/arrow_realtime_store_factory.h b/include/paimon/realtime/arrow_realtime_store_factory.h index 235671ac1..a2e3f22af 100644 --- a/include/paimon/realtime/arrow_realtime_store_factory.h +++ b/include/paimon/realtime/arrow_realtime_store_factory.h @@ -28,6 +28,8 @@ class PAIMON_EXPORT ArrowRealtimeStoreFactory : public RealtimeStoreFactory { public: /// Creates an Arrow-backed append or primary-key store for one partition and bucket. /// The framework uses the append implementation as the delegate of its DEDUPLICATE wrapper. + /// Append delegates resolve configured index plugins through the registry. + /// Primary-key stores do not support query indexes. Result> Create(RealtimeStoreCreateRequest&& request) override; }; diff --git a/include/paimon/realtime/index/realtime_index.h b/include/paimon/realtime/index/realtime_index.h new file mode 100644 index 000000000..0a3717dd8 --- /dev/null +++ b/include/paimon/realtime/index/realtime_index.h @@ -0,0 +1,49 @@ +/* + * 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 "paimon/realtime/offset_range.h" +#include "paimon/result.h" +#include "paimon/visibility.h" + +struct ArrowArray; + +namespace paimon { +class RealtimeIndexReader; + +/// Append-only in-memory index for one field, supporting one writer and concurrent readers. +class PAIMON_EXPORT RealtimeIndex { + public: + virtual ~RealtimeIndex() = default; + + /// Consumes [_REALTIME_OFFSET, field] with increasing offsets. + virtual Status AddBatch(::ArrowArray* batch) = 0; + + /// Creates a reader for [begin, end), sharing postings without copying. + virtual Result> CreateReader( + const OffsetRange& offset_range) const = 0; + + virtual uint64_t GetMemoryUsage() const = 0; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/index/realtime_index_factory.h b/include/paimon/realtime/index/realtime_index_factory.h new file mode 100644 index 000000000..785b2190d --- /dev/null +++ b/include/paimon/realtime/index/realtime_index_factory.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 +#include +#include + +#include "paimon/factories/factory.h" +#include "paimon/result.h" +#include "paimon/visibility.h" + +struct ArrowSchema; + +namespace paimon { +class MemoryPool; +class RealtimeIndex; + +/// Factory for field indexes configured by realtime.index..columns. +class PAIMON_EXPORT RealtimeIndexFactory : public Factory { + public: + ~RealtimeIndexFactory() override = default; + + /// Registry suffix for real-time index plugins. + static const char REALTIME_INDEX_IDENTIFIER_SUFFIX[]; + + /// Resolves identifier + "-realtime" and consumes the [_REALTIME_OFFSET, field] schema. + static Result> Get( + const std::string& identifier, ::ArrowSchema* write_schema, + const std::map& options, const std::shared_ptr& pool); + + /// Creates a new field index; consumes the [_REALTIME_OFFSET, field] schema. + virtual Result> Create( + ::ArrowSchema* write_schema, const std::map& options, + const std::shared_ptr& pool) const = 0; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/index/realtime_index_reader.h b/include/paimon/realtime/index/realtime_index_reader.h new file mode 100644 index 000000000..78e1a3681 --- /dev/null +++ b/include/paimon/realtime/index/realtime_index_reader.h @@ -0,0 +1,56 @@ +/* + * 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 "paimon/predicate/function_visitor.h" +#include "paimon/utils/roaring_bitmap64.h" +#include "paimon/visibility.h" + +namespace paimon { + +/// Visits predicates on one indexed field, returning realtime offsets. +/// nullopt means no pruning; an empty bitmap means no candidates. False negatives are not allowed. +/// Unsupported predicates default to no pruning. Plugins override only the supported visits. +class PAIMON_EXPORT RealtimeIndexReader : public FunctionVisitor> { + public: + ~RealtimeIndexReader() override = default; + + Result> VisitIsNotNull() override; + Result> VisitIsNull() override; + Result> VisitEqual(const Literal& literal) override; + Result> VisitNotEqual(const Literal& literal) override; + Result> VisitLessThan(const Literal& literal) override; + Result> VisitLessOrEqual(const Literal& literal) override; + Result> VisitGreaterThan(const Literal& literal) override; + Result> VisitGreaterOrEqual(const Literal& literal) override; + Result> VisitIn(const std::vector& literals) override; + Result> VisitNotIn( + const std::vector& literals) override; + Result> VisitStartsWith(const Literal& literal) override; + Result> VisitEndsWith(const Literal& literal) override; + Result> VisitContains(const Literal& literal) override; + Result> VisitLike(const Literal& literal) override; +}; + +} // namespace paimon diff --git a/include/paimon/realtime/realtime_context.h b/include/paimon/realtime/realtime_context.h index ca410bb6d..bda184e79 100644 --- a/include/paimon/realtime/realtime_context.h +++ b/include/paimon/realtime/realtime_context.h @@ -88,6 +88,8 @@ class PAIMON_EXPORT RealtimeMetrics { /// reads. `RealtimeContext` itself is not a customization interface and must not be implemented by /// applications. Customize real-time storage and retrieval through `RealtimeStoreFactory` and /// `RealtimeStore` instead. +/// Only one writer may actively use a partition-bucket in this context. Close that writer before +/// creating its replacement; writers for different partition-buckets may share the context. /// /// A context is valid only for one uninterrupted committed-progress history. Overwrite, truncate, /// partition drop, and rollback operations do not automatically clear process-local real-time diff --git a/include/paimon/realtime/realtime_store.h b/include/paimon/realtime/realtime_store.h index da0165845..d1f3ee0b3 100644 --- a/include/paimon/realtime/realtime_store.h +++ b/include/paimon/realtime/realtime_store.h @@ -142,6 +142,9 @@ struct PAIMON_EXPORT RealtimeQueryContext { /// Optional predicate using field indexes from `read_schema`. A non-null predicate allows the /// plugin to prune candidate rows. Exact filtering is applied by the Paimon read framework. std::shared_ptr predicate; + /// Required [disk snapshot committed end, pinned memory end) bounds. + /// Pin memory before selecting the disk snapshot; keep these bounds for the query. + OffsetRange offset_range; /// Maximum number of rows in a returned query batch. int32_t read_batch_size = 1024; }; diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 4a5077c31..461081bba 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -424,6 +424,11 @@ set(PAIMON_CORE_SRCS core/realtime/arrow_realtime_store_factory.cpp core/realtime/data_file_key_offset_index_writer.cpp core/realtime/key_offset_index_utils.cpp + core/realtime/index/hash/hash_realtime_index.cpp + core/realtime/index/realtime_index_evaluator.cpp + core/realtime/index/realtime_index_factory.cpp + core/realtime/index/realtime_index_reader.cpp + core/realtime/index/realtime_index_set.cpp core/realtime/realtime_key_offset_index.cpp core/realtime/realtime_file_schema.cpp core/realtime/realtime_offset_file_index_lookup.cpp @@ -914,6 +919,7 @@ if(PAIMON_BUILD_TESTS) core/realtime/arrow_deduplicate_realtime_store_test.cpp core/realtime/arrow_realtime_store_test.cpp core/realtime/key_offset_index_utils_test.cpp + core/realtime/index/realtime_index_test.cpp core/realtime/primary_key_realtime_store_test.cpp core/realtime/realtime_deduplicate_state_test.cpp core/realtime/realtime_file_schema_test.cpp diff --git a/src/paimon/common/io/offset_output_stream.cpp b/src/paimon/common/io/offset_output_stream.cpp index e5cd922f0..727116c7f 100644 --- a/src/paimon/common/io/offset_output_stream.cpp +++ b/src/paimon/common/io/offset_output_stream.cpp @@ -39,6 +39,14 @@ Result OffsetOutputStream::GetPos() const { return position - start_; } +Result OffsetOutputStream::GetUnderlyingPos() const { + const OutputStream* output = output_; + while (const auto* offset_output = dynamic_cast(output)) { + output = offset_output->output_; + } + return output->GetPos(); +} + Result OffsetOutputStream::GetUri() const { return output_->GetUri(); } diff --git a/src/paimon/common/io/offset_output_stream.h b/src/paimon/common/io/offset_output_stream.h index 3a2410bb1..95a0defb5 100644 --- a/src/paimon/common/io/offset_output_stream.h +++ b/src/paimon/common/io/offset_output_stream.h @@ -39,6 +39,8 @@ class PAIMON_EXPORT OffsetOutputStream : public OutputStream { Result Write(const char* buffer, int64_t size) override; Result GetPos() const override; + /// Returns the current position of the innermost stream, unwrapping nested OffsetOutputStreams. + Result GetUnderlyingPos() const; Result GetUri() const override; Status Flush() override; Status Close() override; diff --git a/src/paimon/common/io/offset_output_stream_test.cpp b/src/paimon/common/io/offset_output_stream_test.cpp index d4ef8c14a..b267bb281 100644 --- a/src/paimon/common/io/offset_output_stream_test.cpp +++ b/src/paimon/common/io/offset_output_stream_test.cpp @@ -62,6 +62,48 @@ TEST(OffsetOutputStreamTest, TestRelativePositionsAndWrites) { ASSERT_EQ("prefixfirstsecond", std::string(bytes->data(), bytes->size())); } +TEST(OffsetOutputStreamTest, TestUnderlyingPositionThroughNestedViews) { + auto pool = GetDefaultPool(); + auto output = + std::make_unique(std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool)); + ASSERT_OK_AND_ASSIGN(int64_t written, output->Write("prefix", 6)); + ASSERT_EQ(6, written); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first, + OffsetOutputStream::Create(output.get())); + ASSERT_OK_AND_ASSIGN(int64_t position, first->GetUnderlyingPos()); + ASSERT_EQ(6, position); + ASSERT_OK_AND_ASSIGN(written, first->Write("first", 5)); + ASSERT_EQ(5, written); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second, + OffsetOutputStream::Create(first.get())); + ASSERT_OK_AND_ASSIGN(position, second->GetPos()); + ASSERT_EQ(0, position); + ASSERT_OK_AND_ASSIGN(position, second->GetUnderlyingPos()); + ASSERT_EQ(11, position); + ASSERT_OK_AND_ASSIGN(written, second->Write("second", 6)); + ASSERT_EQ(6, written); + ASSERT_OK_AND_ASSIGN(position, first->GetPos()); + ASSERT_EQ(11, position); + ASSERT_OK_AND_ASSIGN(position, second->GetPos()); + ASSERT_EQ(6, position); + ASSERT_OK_AND_ASSIGN(position, second->GetUnderlyingPos()); + ASSERT_EQ(17, position); + + // Direct writes to the underlying stream must be reflected by every view. + ASSERT_OK_AND_ASSIGN(written, output->Write("suffix", 6)); + ASSERT_EQ(6, written); + ASSERT_OK_AND_ASSIGN(position, first->GetUnderlyingPos()); + ASSERT_EQ(23, position); + ASSERT_OK_AND_ASSIGN(position, second->GetUnderlyingPos()); + ASSERT_EQ(23, position); + ASSERT_OK_AND_ASSIGN(position, second->GetPos()); + ASSERT_EQ(12, position); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR bytes, output->Finish(pool.get())); + ASSERT_EQ("prefixfirstsecondsuffix", std::string(bytes->data(), bytes->size())); +} + TEST(OffsetOutputStreamTest, TestFlushAndClose) { auto dir = UniqueTestDirectory::Create(); auto fs = dir->GetFileSystem(); diff --git a/src/paimon/common/reader/reader_utils.cpp b/src/paimon/common/reader/reader_utils.cpp index 0fe38d2c3..a0bbaf5c7 100644 --- a/src/paimon/common/reader/reader_utils.cpp +++ b/src/paimon/common/reader/reader_utils.cpp @@ -41,6 +41,9 @@ Result ReaderUtils::GenerateFilteredArrayVector( if (bitmap.Cardinality() == 0) { return Status::Invalid("selection bitmap cannot be empty in GenerateFilteredArrayVector"); } + if (bitmap.Cardinality() == src_array->length()) { + return arrow::ArrayVector{src_array}; + } arrow::ArrayVector array_vec; auto valid_iter = bitmap.Begin(); int32_t pos = 0; diff --git a/src/paimon/core/operation/abstract_file_store_write.cpp b/src/paimon/core/operation/abstract_file_store_write.cpp index 3c8e19114..78a2ac31c 100644 --- a/src/paimon/core/operation/abstract_file_store_write.cpp +++ b/src/paimon/core/operation/abstract_file_store_write.cpp @@ -372,6 +372,41 @@ Result> AbstractFileStoreWrite::PrepareCommi return result; } +Status AbstractFileStoreWrite::RefreshRealtimeCommittedProgress( + int64_t snapshot_id, const RealtimeOffsetMap& committed_offsets, + const RealtimeDataFileMap& committed_data_files) { + // Prepare drains each bucket's pending queue into local handles. Keep it and Seal out until + // the snapshot is installed and the covered handles have been removed from these queues. + std::lock_guard seal_prepare_lock(realtime_seal_prepare_mutex_); + std::vector, int64_t>> writers_to_release; + { + std::lock_guard lock(writers_mutex_); + for (const auto& [partition, buckets] : writers_) { + std::vector> partition_values; + PAIMON_ASSIGN_OR_RAISE(partition_values, + file_store_path_factory_->GeneratePartitionVector(partition)); + std::map partition_map(partition_values.begin(), + partition_values.end()); + for (const auto& [bucket, container] : buckets) { + auto iter = committed_offsets.find(RealtimePartitionBucket(partition_map, bucket)); + if (iter != committed_offsets.end()) { + writers_to_release.emplace_back(container.writer, iter->second); + } + } + } + } + // No writer-registry lock is held during snapshot preparation I/O, so writes can continue + // and register new bucket writers. Those writers cannot seal until this refresh completes. + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr context, + RealtimeContextImpl::Cast(GetRealtimeContext())); + PAIMON_RETURN_NOT_OK( + context->AdvanceCommittedProgress(snapshot_id, committed_offsets, committed_data_files)); + for (const auto& [writer, committed_end_offset] : writers_to_release) { + writer->ReleaseCommittedRealtimeSegments(committed_end_offset); + } + return Status::OK(); +} + Result> AbstractFileStoreWrite::PrepareRealtimeCommit() { struct WriterSnapshot { BinaryRow partition; diff --git a/src/paimon/core/operation/abstract_file_store_write.h b/src/paimon/core/operation/abstract_file_store_write.h index 644d88dd7..8ed3d5266 100644 --- a/src/paimon/core/operation/abstract_file_store_write.h +++ b/src/paimon/core/operation/abstract_file_store_write.h @@ -37,6 +37,7 @@ #include "paimon/core/deletionvectors/bucketed_dv_maintainer.h" #include "paimon/core/index/pk/bucketed_primary_key_index_maintainer.h" #include "paimon/core/memory/writer_memory_manager.h" +#include "paimon/core/realtime/realtime_context_impl.h" #include "paimon/file_store_write.h" #include "paimon/logging.h" #include "paimon/metrics.h" @@ -147,6 +148,10 @@ class AbstractFileStoreWrite : public FileStoreWrite { return nullptr; } + Status RefreshRealtimeCommittedProgress(int64_t snapshot_id, + const RealtimeOffsetMap& committed_offsets, + const RealtimeDataFileMap& committed_data_files); + Result> ScanExistingFileMetas(const BinaryRow& partition, int32_t bucket) const; int32_t GetDefaultBucketNum() const; diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index 666260267..5887ddd0b 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -152,11 +152,7 @@ Status AppendOnlyFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { committed_data_files[RealtimePartitionBucket(std::move(partition), entry.Bucket())] .push_back(entry.File()); } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(realtime_context_)); - PAIMON_RETURN_NOT_OK(realtime_context_impl->AdvanceCommittedProgress( - snapshot_id, committed_offsets, committed_data_files)); - return Status::OK(); + return RefreshRealtimeCommittedProgress(snapshot_id, committed_offsets, committed_data_files); } Result> AppendOnlyFileStoreWrite::CreateFileStoreScan( diff --git a/src/paimon/core/operation/key_value_file_store_write.cpp b/src/paimon/core/operation/key_value_file_store_write.cpp index be72ad914..830e6c531 100644 --- a/src/paimon/core/operation/key_value_file_store_write.cpp +++ b/src/paimon/core/operation/key_value_file_store_write.cpp @@ -193,10 +193,8 @@ Status KeyValueFileStoreWrite::RefreshCommittedSnapshot(int64_t snapshot_id) { RealtimeOffsetMap committed_offsets, RealtimeCommitProperties::ReadOffsets(std::optional(std::move(snapshot)), options_.GetFileSystem())); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_context_impl, - RealtimeContextImpl::Cast(realtime_context_)); - return realtime_context_impl->AdvanceCommittedProgress(snapshot_id, committed_offsets, - /*committed_data_files=*/{}); + return RefreshRealtimeCommittedProgress(snapshot_id, committed_offsets, + /*committed_data_files=*/{}); } Status KeyValueFileStoreWrite::Close() { diff --git a/src/paimon/core/operation/key_value_file_store_write_test.cpp b/src/paimon/core/operation/key_value_file_store_write_test.cpp index c0e94e236..7602a2763 100644 --- a/src/paimon/core/operation/key_value_file_store_write_test.cpp +++ b/src/paimon/core/operation/key_value_file_store_write_test.cpp @@ -270,7 +270,10 @@ class KeyValueFileStoreWriteTest : public ::testing::Test { schema_layout->StoreWriteSchema(); auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*store_write_schema, c_schema.get())); - RealtimeQueryContext query_context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext query_context{ + c_schema.get(), + /*predicate=*/nullptr, + views[0].read_view->GetOffsetRange().value_or(OffsetRange(0, 0))}; PAIMON_ASSIGN_OR_RAISE( std::vector> readers, views[0].store->CreateQueryReaders(views[0].read_view, query_context)); diff --git a/src/paimon/core/realtime/arrow_deduplicate_realtime_store.cpp b/src/paimon/core/realtime/arrow_deduplicate_realtime_store.cpp index 0607c142b..ebae50976 100644 --- a/src/paimon/core/realtime/arrow_deduplicate_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_deduplicate_realtime_store.cpp @@ -516,7 +516,8 @@ Result>> ArrowDeduplicateRealtimeStore: ArrowSchema inner_c_schema; ArrowSchemaMarkReleased(&inner_c_schema); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*output_schema, &inner_c_schema)); - RealtimeQueryContext inner_context{&inner_c_schema, context.predicate, context.read_batch_size}; + RealtimeQueryContext inner_context{&inner_c_schema, context.predicate, context.offset_range, + context.read_batch_size}; PAIMON_ASSIGN_OR_RAISE( std::vector> inner, delegate_->CreateQueryReaders(deduplicate_view->Delegate(), inner_context)); diff --git a/src/paimon/core/realtime/arrow_deduplicate_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_deduplicate_realtime_store_test.cpp index f064bce35..8cc962044 100644 --- a/src/paimon/core/realtime/arrow_deduplicate_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_deduplicate_realtime_store_test.cpp @@ -40,6 +40,7 @@ #include "paimon/fs/local/local_file_system.h" #include "paimon/macros.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/testing/utils/blocking_file_system.h" #include "paimon/testing/utils/test_helper.h" #include "paimon/testing/utils/testharness.h" @@ -141,10 +142,14 @@ class ArrowDeduplicateRealtimeStoreTest : public testing::Test { }); pool_ = GetDefaultPool(); arrow_pool_ = GetArrowPool(pool_); - delegate_ = std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, - /*temp_directory=*/"", /*spill_file_system=*/nullptr, - /*spill_compression=*/"zstd", /*spill_compression_level=*/1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN(delegate_, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, + /*temp_directory=*/"", + /*spill_file_system=*/nullptr, + /*spill_compression=*/"zstd", + /*spill_compression_level=*/1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create( schema_, {"id"}, delegate_, pool_, arrow_pool_)); empty_lookup_ = std::make_shared(); @@ -173,10 +178,14 @@ class ArrowDeduplicateRealtimeStoreTest : public testing::Test { } void ResetStoreWithBitmap(const RoaringBitmap32& bitmap) { - delegate_ = std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, - /*temp_directory=*/"", /*spill_file_system=*/nullptr, - /*spill_compression=*/"zstd", /*spill_compression_level=*/1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN(delegate_, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, + /*temp_directory=*/"", + /*spill_file_system=*/nullptr, + /*spill_compression=*/"zstd", + /*spill_compression_level=*/1, + /*index_options=*/{}, pool_, arrow_pool_)); auto selecting = std::make_shared(delegate_, bitmap); ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create( schema_, {"id"}, selecting, pool_, arrow_pool_)); @@ -188,7 +197,9 @@ class ArrowDeduplicateRealtimeStoreTest : public testing::Test { int32_t batch_size) const { auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, c_schema.get())); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, batch_size}; + RealtimeQueryContext context{ + c_schema.get(), + /*predicate=*/nullptr, view->GetOffsetRange().value_or(OffsetRange(0, 0)), batch_size}; return store_->CreateQueryReaders(view, context); } @@ -200,9 +211,15 @@ class ArrowDeduplicateRealtimeStoreTest : public testing::Test { } Result> ReadOffsets(const std::shared_ptr& view) const { + return ReadOffsets(view, nullptr, view->GetOffsetRange().value_or(OffsetRange(0, 0))); + } + + Result> ReadOffsets(const std::shared_ptr& view, + const std::shared_ptr& predicate, + const OffsetRange& offset_range) const { auto c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, c_schema.get())); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), predicate, offset_range, 1024}; PAIMON_ASSIGN_OR_RAISE(std::vector> readers, store_->CreateQueryReaders(view, context)); std::vector result; @@ -315,6 +332,46 @@ TEST_F(ArrowDeduplicateRealtimeStoreTest, TestLastWriteWinsAndReadViewIsSnapshot ASSERT_EQ(OffsetRange(1, 5), segment.value()->GetOffsetRange()); } +TEST_F(ArrowDeduplicateRealtimeStoreTest, TestHashPostingsCombineWithPinnedDvWithoutIndexCopies) { + ASSERT_OK_AND_ASSIGN( + delegate_, ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, "", nullptr, "zstd", 1, + {{"realtime.index.hash.columns", "id,value"}}, pool_, + arrow_pool_)); + ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create(schema_, {"id"}, delegate_, + pool_, arrow_pool_)); + ASSERT_OK(store_->AttachCommittedKeyOffsetLookup(empty_lookup_)); + ASSERT_OK(store_->WriteAndLookup(MakeBatch(R"([[10, 1, "old"]])"), + {RecordBatch::RowKind::INSERT}, OffsetRange(10, 11), + RoaringBitmap64())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr old_view, store_->AcquireReadView()); + ASSERT_OK(store_->WriteAndLookup(MakeBatch(R"([[20, 1, "new"]])"), + {RecordBatch::RowKind::UPDATE_AFTER}, OffsetRange(20, 21), + RoaringBitmap64())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr new_view, store_->AcquireReadView()); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store_->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK(store_->RetireCommittedOffsetAndLookup(21, empty_lookup_)); + + const std::shared_ptr key = + PredicateBuilder::Equal(1, "id", FieldType::BIGINT, Literal(int64_t{1})); + const std::shared_ptr old_value = PredicateBuilder::Equal( + 2, "value", FieldType::STRING, Literal(FieldType::STRING, "old", 3)); + ASSERT_OK_AND_ASSIGN(std::vector old_offsets, + ReadOffsets(old_view, key, OffsetRange(0, 100))); + ASSERT_EQ((std::vector{10}), old_offsets); + ASSERT_OK_AND_ASSIGN(std::vector new_offsets, + ReadOffsets(new_view, key, OffsetRange(0, 100))); + ASSERT_EQ((std::vector{20}), new_offsets); + ASSERT_OK_AND_ASSIGN(old_offsets, ReadOffsets(old_view, old_value, OffsetRange(0, 100))); + ASSERT_EQ((std::vector{10}), old_offsets); + ASSERT_OK_AND_ASSIGN(new_offsets, ReadOffsets(new_view, old_value, OffsetRange(0, 100))); + ASSERT_TRUE(new_offsets.empty()); + ASSERT_OK_AND_ASSIGN(new_offsets, ReadOffsets(new_view, key, OffsetRange(10, 20))); + ASSERT_TRUE(new_offsets.empty()); +} + TEST_F(ArrowDeduplicateRealtimeStoreTest, TestDeleteProducesNoLiveRows) { ASSERT_OK(store_->WriteAndLookup(MakeBatch(R"([[1, 1, "value"]])"), {RecordBatch::RowKind::INSERT}, OffsetRange(1, 2), @@ -450,9 +507,10 @@ TEST_F(ArrowDeduplicateRealtimeStoreTest, TestSpilledBitmapFilteringPreservesRaw auto directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); auto fs = std::make_shared(); - delegate_ = std::make_shared(schema_, RealtimeStoreMode::APPEND_ONLY, - StatisticsMode::NONE, directory->Str(), fs, - "zstd", 1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + delegate_, ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, directory->Str(), fs, "zstd", 1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create(schema_, {"id"}, delegate_, pool_, arrow_pool_)); ASSERT_OK(store_->AttachCommittedKeyOffsetLookup(empty_lookup_)); @@ -557,9 +615,11 @@ TEST_F(ArrowDeduplicateRealtimeStoreTest, TestOffsetMergeScanMatchesContains) { TEST_F(ArrowDeduplicateRealtimeStoreTest, TestFreezePreservesCheckpointWithoutSpilling) { auto directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); - delegate_ = std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, directory->Str(), - directory->GetFileSystem(), "zstd", 1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + delegate_, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, + directory->Str(), directory->GetFileSystem(), "zstd", 1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create(schema_, {"id"}, delegate_, pool_, arrow_pool_)); ASSERT_OK(store_->AttachCommittedKeyOffsetLookup(empty_lookup_)); @@ -642,9 +702,10 @@ TEST_F(ArrowDeduplicateRealtimeStoreTest, TestSpillDoesNotBlockWriteAndPreserves auto close_gate = fs->BlockNext(BlockingFileSystem::Operation::CLOSE, /*path=*/""); close_gate->Release(Status::IOError("injected spill close failure")); } - delegate_ = std::make_shared(schema_, RealtimeStoreMode::APPEND_ONLY, - StatisticsMode::NONE, directory->Str(), fs, - "zstd", 1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN(delegate_, ArrowRealtimeStore::Create( + schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, directory->Str(), fs, "zstd", 1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_OK_AND_ASSIGN(store_, ArrowDeduplicateRealtimeStore::Create( schema_, {"id"}, delegate_, pool_, arrow_pool_)); ASSERT_OK(store_->AttachCommittedKeyOffsetLookup(empty_lookup_)); diff --git a/src/paimon/core/realtime/arrow_realtime_store.cpp b/src/paimon/core/realtime/arrow_realtime_store.cpp index 7f9212741..ad4ea6d72 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store.cpp @@ -26,11 +26,15 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "arrow/compute/api.h" #include "arrow/compute/api_aggregate.h" #include "paimon/common/data/columnar/columnar_array.h" #include "paimon/common/data/columnar/columnar_row.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/table/special_fields.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" @@ -40,6 +44,9 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/core/disk/io_manager.h" #include "paimon/core/io/arrow_ipc_file.h" +#include "paimon/core/realtime/index/realtime_index_evaluator.h" +#include "paimon/core/realtime/index/realtime_index_set.h" +#include "paimon/core/realtime/realtime_utils.h" #include "paimon/core/utils/nested_projection_utils.h" #include "paimon/fs/file_system.h" #include "paimon/macros.h" @@ -109,28 +116,50 @@ class SlicedQueryBatchReader final : public BatchReader { : reader_(std::move(reader)), read_batch_size_(read_batch_size), arrow_pool_(arrow_pool) {} Result NextBatch() override { - while (!current_batch_ || next_row_ >= current_batch_->length()) { - current_batch_.reset(); - next_row_ = 0; - PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, reader_->NextBatch()); - if (IsEofBatch(batch)) { - return batch; + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, NextBatchWithBitmap()); + return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), arrow_pool_); + } + + Result NextBatchWithBitmap() override { + while (true) { + if (!current_batch_ || next_row_ >= current_batch_->length()) { + current_batch_.reset(); + next_row_ = 0; + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, reader_->NextBatchWithBitmap()); + if (IsEofBatch(batch)) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ImportArray(batch.first.first.get(), batch.first.second.get())); + if (!array || array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("real-time query batch must be a StructArray"); + } + current_batch_ = checked_pointer_cast(array); + current_bitmap_ = std::move(batch.second); } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr array, - arrow::ImportArray(batch.first.get(), batch.second.get())); - if (!array || array->type_id() != arrow::Type::STRUCT) { - return Status::Invalid("real-time query batch must be a StructArray"); + + const int64_t length = + std::min(read_batch_size_, current_batch_->length() - next_row_); + const int64_t begin = next_row_; + next_row_ += length; + RoaringBitmap32 selection; + if (begin == 0 && length == current_batch_->length()) { + selection = std::move(current_bitmap_); + } else { + for (auto iter = current_bitmap_.EqualOrLarger(begin); + iter != current_bitmap_.End() && *iter < next_row_; ++iter) { + selection.Add(*iter - begin); + } } - current_batch_ = checked_pointer_cast(array); + if (selection.IsEmpty()) { + continue; + } + std::shared_ptr slice = + checked_pointer_cast(current_batch_->Slice(begin, length)); + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, ExportStructBatch(slice, arrow_pool_)); + return ReadBatchWithBitmap(std::move(batch), std::move(selection)); } - - const int64_t length = - std::min(read_batch_size_, current_batch_->length() - next_row_); - std::shared_ptr slice = - checked_pointer_cast(current_batch_->Slice(next_row_, length)); - next_row_ += length; - return ExportStructBatch(slice, arrow_pool_); } std::shared_ptr GetReaderMetrics() const override { @@ -139,6 +168,7 @@ class SlicedQueryBatchReader final : public BatchReader { void Close() override { current_batch_.reset(); + current_bitmap_ = RoaringBitmap32(); arrow_pool_.reset(); reader_->Close(); } @@ -148,6 +178,7 @@ class SlicedQueryBatchReader final : public BatchReader { int32_t read_batch_size_; std::shared_ptr arrow_pool_; std::shared_ptr current_batch_; + RoaringBitmap32 current_bitmap_; int64_t next_row_ = 0; }; @@ -165,8 +196,9 @@ std::vector> SliceQueryReaders( class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { public: - Segment(const OffsetRange& offset_range, int64_t row_count) - : offset_range_(offset_range), row_count_(row_count) {} + Segment(const OffsetRange& offset_range, int64_t row_count, + std::shared_ptr index) + : offset_range_(offset_range), row_count_(row_count), index_(std::move(index)) {} OffsetRange GetOffsetRange() const override { return offset_range_; @@ -178,15 +210,22 @@ class ArrowRealtimeStore::Segment : public RealtimeSegmentHandle { virtual uint64_t GetMemoryUsage() const = 0; + const std::shared_ptr& GetIndex() const { + return index_; + } + private: OffsetRange offset_range_; int64_t row_count_; + std::shared_ptr index_; }; class ArrowRealtimeStore::MemorySegment final : public ArrowRealtimeStore::Segment { public: - MemorySegment(const OffsetRange& offset_range, std::vector&& batches) - : Segment(offset_range, CountRows(batches)), batches_(std::move(batches)) { + MemorySegment(const OffsetRange& offset_range, std::vector&& batches, + std::shared_ptr index) + : Segment(offset_range, CountRows(batches), std::move(index)), + batches_(std::move(batches)) { for (const StoredBatch& batch : batches_) { memory_usage_ += batch.memory_usage; } @@ -197,7 +236,7 @@ class ArrowRealtimeStore::MemorySegment final : public ArrowRealtimeStore::Segme } uint64_t GetMemoryUsage() const override { - return memory_usage_; + return memory_usage_ + (GetIndex() ? GetIndex()->GetMemoryUsage() : 0); } private: @@ -243,8 +282,11 @@ class ArrowRealtimeStore::SpillFile { class ArrowRealtimeStore::SpilledSegment final : public ArrowRealtimeStore::Segment { public: SpilledSegment(const OffsetRange& offset_range, int64_t row_count, - std::vector&& batches, std::shared_ptr file) - : Segment(offset_range, row_count), batches_(std::move(batches)), file_(std::move(file)) { + std::vector&& batches, std::shared_ptr file, + std::shared_ptr index) + : Segment(offset_range, row_count, std::move(index)), + batches_(std::move(batches)), + file_(std::move(file)) { for (const SpilledBatch& batch : batches_) { memory_usage_ += batch.memory_usage; } @@ -259,7 +301,7 @@ class ArrowRealtimeStore::SpilledSegment final : public ArrowRealtimeStore::Segm } uint64_t GetMemoryUsage() const override { - return memory_usage_; + return memory_usage_ + (GetIndex() ? GetIndex()->GetMemoryUsage() : 0); } private: @@ -441,6 +483,8 @@ class ArrowRealtimeStore::AppendQueryBatchReader : public BatchReader { const std::shared_ptr& read_schema, const std::shared_ptr& predicate_filter, std::vector&& statistics_mapping, + const std::shared_ptr& predicate, + const OffsetRange& offset_range, const std::shared_ptr& arrow_pool, const std::shared_ptr& memory_pool) : view_(view), @@ -449,73 +493,36 @@ class ArrowRealtimeStore::AppendQueryBatchReader : public BatchReader { memory_pool_(memory_pool), predicate_filter_(predicate_filter), statistics_mapping_(std::move(statistics_mapping)), + predicate_(predicate), + offset_range_(offset_range), metrics_(std::make_shared()) {} Result NextBatch() override { + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch, NextBatchWithBitmap()); + return ReaderUtils::ApplyBitmapToReadBatch(std::move(batch), arrow_pool_); + } + + Result NextBatchWithBitmap() override { while (view_ && next_segment_ < view_->GetSegments().size()) { const std::shared_ptr& segment = view_->GetSegments()[next_segment_]; - std::shared_ptr memory_segment = - std::dynamic_pointer_cast(segment); - std::shared_ptr spilled_segment = - std::dynamic_pointer_cast(segment); - if (!memory_segment && !spilled_segment) { - return Status::Invalid("unknown Arrow real-time segment type"); - } - size_t batch_count = memory_segment ? memory_segment->GetBatches().size() - : spilled_segment->GetBatches().size(); - if (next_batch_ >= batch_count) { - ++next_segment_; - next_batch_ = 0; - spill_reader_.reset(); + PAIMON_ASSIGN_OR_RAISE(bool may_match, PrepareSegment(segment)); + if (!may_match) { + NextSegment(); continue; } - int64_t row_count; - const std::optional* statistics; - std::shared_ptr source; - if (memory_segment) { - const StoredBatch& stored = memory_segment->GetBatches()[next_batch_]; - row_count = stored.data->length(); - statistics = &stored.statistics; - source = stored.data; - } else if (spilled_segment) { - const SpilledBatch& stored = spilled_segment->GetBatches()[next_batch_]; - row_count = stored.row_count; - statistics = &stored.statistics; - } else { - return Status::Invalid("unknown Arrow real-time segment type"); - } - PAIMON_ASSIGN_OR_RAISE(bool may_match, - ArrowRealtimeStore::MayMatchStatistics( - row_count, *statistics, read_schema_, predicate_filter_, - statistics_mapping_, memory_pool_)); - if (!may_match) { - ++next_batch_; + PAIMON_ASSIGN_OR_RAISE(bool has_batch, SeekNextBatch(segment)); + if (!has_batch) { + NextSegment(); continue; } - if (spilled_segment) { - if (!spill_reader_) { - PAIMON_ASSIGN_OR_RAISE( - spill_reader_, - ArrowIpcFileReader::Open(spilled_segment->GetFile()->GetFileSystem(), - spilled_segment->GetFile()->GetPath(), - /*use_threads=*/false, arrow_pool_)); - if (spill_reader_->GetRecordBatchCount() != - static_cast(spilled_segment->GetBatches().size())) { - return Status::Invalid( - "real-time spill file batch count does not match segment metadata"); - } - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr record_batch, - spill_reader_->ReadRecordBatch(next_batch_)); - PAIMON_ASSIGN_OR_RAISE(source, ToStructArray(record_batch)); + const size_t batch_index = next_batch_++; + PAIMON_ASSIGN_OR_RAISE(std::optional batch, + TryReadBatch(segment, batch_index)); + if (batch) { + return std::move(*batch); } - ++next_batch_; - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, - ProjectBatch(source, read_schema_, arrow_pool_.get())); - PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, ExportStructBatch(output, arrow_pool_)); - return batch; } - return MakeEofBatch(); + return MakeEofBatchWithBitmap(); } std::shared_ptr GetReaderMetrics() const override { @@ -525,27 +532,249 @@ class ArrowRealtimeStore::AppendQueryBatchReader : public BatchReader { void Close() override { view_ = nullptr; spill_reader_.reset(); + candidates_.reset(); } private: + Result PrepareSegment(const std::shared_ptr& segment) { + const OffsetRange segment_range = segment->GetOffsetRange(); + if (segment_range.end <= offset_range_.begin || segment_range.begin >= offset_range_.end) { + return false; + } + if (!segment_initialized_) { + if (segment->GetIndex()) { + PAIMON_ASSIGN_OR_RAISE( + candidates_, RealtimeIndexEvaluator::Evaluate( + segment->GetIndex(), predicate_, + OffsetRange(std::max(segment_range.begin, offset_range_.begin), + std::min(segment_range.end, offset_range_.end)))); + } + segment_initialized_ = true; + } + return !candidates_ || !candidates_->IsEmpty(); + } + + Result SeekNextBatch(const std::shared_ptr& segment) { + const std::shared_ptr memory_segment = + std::dynamic_pointer_cast(segment); + if (memory_segment) { + return SeekNextBatch(memory_segment->GetBatches()); + } + const std::shared_ptr spilled_segment = + std::dynamic_pointer_cast(segment); + if (spilled_segment) { + return SeekNextBatch(spilled_segment->GetBatches()); + } + return Status::Invalid("unknown Arrow real-time segment type"); + } + + template + bool SeekNextBatch(const std::vector& batches) { + if (next_batch_ >= batches.size()) { + return false; + } + const OffsetRange current_range = batches[next_batch_].offset_range; + int64_t first_offset = std::max(current_range.begin, offset_range_.begin); + if (candidates_) { + const RoaringBitmap64::Iterator first = candidates_->EqualOrLarger(first_offset); + if (first == candidates_->End() || *first >= offset_range_.end) { + return false; + } + first_offset = *first; + } + const auto next = std::lower_bound( + batches.begin() + next_batch_, batches.end(), first_offset, + [](const Batch& batch, int64_t offset) { return batch.offset_range.end <= offset; }); + next_batch_ = next - batches.begin(); + return next_batch_ < batches.size(); + } + + Result> TryReadBatch(const std::shared_ptr& segment, + size_t batch_index) { + const std::shared_ptr memory_segment = + std::dynamic_pointer_cast(segment); + const std::shared_ptr spilled_segment = + std::dynamic_pointer_cast(segment); + int64_t row_count; + OffsetRange batch_range(0, 0); + const std::optional* statistics; + std::shared_ptr source; + if (memory_segment) { + const StoredBatch& stored = memory_segment->GetBatches()[batch_index]; + row_count = stored.data->length(); + statistics = &stored.statistics; + source = stored.data; + batch_range = stored.offset_range; + } else if (spilled_segment) { + const SpilledBatch& stored = spilled_segment->GetBatches()[batch_index]; + row_count = stored.row_count; + statistics = &stored.statistics; + batch_range = stored.offset_range; + } else { + return Status::Invalid("unknown Arrow real-time segment type"); + } + const OffsetRange visible_range(std::max(batch_range.begin, offset_range_.begin), + std::min(batch_range.end, offset_range_.end)); + if (visible_range.begin >= visible_range.end || + (candidates_ && !candidates_->ContainsAny(visible_range.begin, visible_range.end))) { + return std::optional(); + } + PAIMON_ASSIGN_OR_RAISE(bool may_match, + ArrowRealtimeStore::MayMatchStatistics( + row_count, *statistics, read_schema_, predicate_filter_, + statistics_mapping_, memory_pool_)); + if (!may_match) { + return std::optional(); + } + if (spilled_segment) { + PAIMON_ASSIGN_OR_RAISE(source, ReadSpilledBatch(spilled_segment, batch_index)); + } + RoaringBitmap32 selection = SelectRows(source, visible_range); + if (selection.IsEmpty()) { + return std::optional(); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr output, + ProjectBatch(source, read_schema_, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, ExportStructBatch(output, arrow_pool_)); + return std::optional( + ReadBatchWithBitmap(std::move(batch), std::move(selection))); + } + + Result> ReadSpilledBatch( + const std::shared_ptr& segment, size_t batch_index) { + if (!spill_reader_) { + PAIMON_ASSIGN_OR_RAISE(spill_reader_, + ArrowIpcFileReader::Open(segment->GetFile()->GetFileSystem(), + segment->GetFile()->GetPath(), + /*use_threads=*/false, arrow_pool_)); + if (spill_reader_->GetRecordBatchCount() != + static_cast(segment->GetBatches().size())) { + return Status::Invalid( + "real-time spill file batch count does not match segment metadata"); + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr record_batch, + spill_reader_->ReadRecordBatch(batch_index)); + return ToStructArray(record_batch); + } + + void NextSegment() { + ++next_segment_; + next_batch_ = 0; + spill_reader_.reset(); + candidates_.reset(); + segment_initialized_ = false; + } + + RoaringBitmap32 SelectRows(const std::shared_ptr& source, + const OffsetRange& visible_range) const { + // TODO(xinyu.lxy): Push down the real-time DV to avoid another offset-to-row mapping + // upstream. + RoaringBitmap32 selection; + if (source->length() == 0) { + return selection; + } + const std::shared_ptr offsets = checked_pointer_cast( + source->GetFieldByName(SpecialFields::RealtimeOffset().Name())); + const int64_t* values = offsets->raw_values(); + const int64_t* begin = + std::lower_bound(values, values + offsets->length(), visible_range.begin); + const int64_t* end = std::lower_bound(begin, values + offsets->length(), visible_range.end); + if (!candidates_) { + selection.AddRange(begin - values, end - values); + return selection; + } + // Offsets may have gaps: sparse candidates seek by offset; dense candidates scan. + const RoaringBitmap64::Iterator sentinel = candidates_->End(); + const RoaringBitmap64::Iterator first = candidates_->EqualOrLarger(visible_range.begin); + RoaringBitmap64::Iterator probe = first; + const int64_t dense_threshold = std::max(1, (end - begin) / 4); + int64_t candidate_count = 0; + while (candidate_count < dense_threshold && probe != sentinel && + *probe < visible_range.end) { + ++candidate_count; + ++probe; + } + if (candidate_count < dense_threshold) { + for (RoaringBitmap64::Iterator iter = first; + iter != sentinel && *iter < visible_range.end; ++iter) { + begin = std::lower_bound(begin, end, *iter); + if (begin != end && *begin == *iter) { + selection.Add(begin - values); + } + } + } else { + for (const int64_t* row = begin; row != end; ++row) { + if (candidates_->Contains(*row)) { + selection.Add(row - values); + } + } + } + return selection; + } + std::shared_ptr view_; std::shared_ptr read_schema_; std::shared_ptr arrow_pool_; std::shared_ptr memory_pool_; std::shared_ptr predicate_filter_; std::vector statistics_mapping_; + std::shared_ptr predicate_; + OffsetRange offset_range_; + std::optional candidates_; + bool segment_initialized_ = false; std::shared_ptr metrics_; std::unique_ptr spill_reader_; size_t next_segment_ = 0; size_t next_batch_ = 0; }; +Result> ArrowRealtimeStore::Create( + const std::shared_ptr& write_schema, RealtimeStoreMode mode, + StatisticsMode statistics_mode, const std::string& temp_directory, + const std::shared_ptr& spill_file_system, const std::string& spill_compression, + int32_t spill_compression_level, const std::map& index_options, + const std::shared_ptr& memory_pool, + const std::shared_ptr& arrow_pool) { + if (!write_schema || !memory_pool || !arrow_pool) { + return Status::Invalid("real-time store requires a schema and memory pools"); + } + if (mode != RealtimeStoreMode::APPEND_ONLY && mode != RealtimeStoreMode::PRIMARY_KEY) { + return Status::Invalid("invalid real-time store mode: ", static_cast(mode)); + } + PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(write_schema)); + if (!temp_directory.empty() && !spill_file_system) { + return Status::Invalid("real-time spill requires a file system"); + } + const bool has_index = RealtimeIndexSet::IsEnabled(index_options); + if (has_index && mode != RealtimeStoreMode::APPEND_ONLY) { + return Status::NotImplemented("real-time query indexes require an append-store delegate"); + } + std::shared_ptr io_manager; + if (!temp_directory.empty()) { + io_manager = std::make_shared(temp_directory, spill_file_system); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr building_index, + RealtimeIndexSet::Create(write_schema, index_options, memory_pool)); + auto store = std::shared_ptr( + new ArrowRealtimeStore(write_schema, mode, statistics_mode, spill_file_system, + spill_compression, spill_compression_level, index_options, + io_manager, std::move(building_index), memory_pool, arrow_pool)); + return store; +} + +Result> ArrowRealtimeStore::CreateBuildingIndex() const { + return RealtimeIndexSet::Create(write_schema_, index_options_, memory_pool_); +} + ArrowRealtimeStore::ArrowRealtimeStore(const std::shared_ptr& write_schema, RealtimeStoreMode mode, StatisticsMode statistics_mode, - const std::string& temp_directory, const std::shared_ptr& spill_file_system, const std::string& spill_compression, int32_t spill_compression_level, + const std::map& index_options, + const std::shared_ptr& io_manager, + std::shared_ptr building_index, const std::shared_ptr& memory_pool, const std::shared_ptr& arrow_pool) : write_schema_(write_schema), @@ -553,13 +782,12 @@ ArrowRealtimeStore::ArrowRealtimeStore(const std::shared_ptr& wri arrow_pool_(arrow_pool), mode_(mode), statistics_mode_(statistics_mode), + io_manager_(io_manager), spill_file_system_(spill_file_system), spill_compression_(spill_compression), - spill_compression_level_(spill_compression_level) { - if (!temp_directory.empty()) { - io_manager_ = std::make_shared(temp_directory, spill_file_system_); - } -} + spill_compression_level_(spill_compression_level), + index_options_(index_options), + building_index_(std::move(building_index)) {} Result> ArrowRealtimeStore::CollectStatistics( const std::shared_ptr& data) const { @@ -671,6 +899,13 @@ 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 ordered and non-overlapping"); } + if (building_index_) { + ArrowArray index_batch; + ArrowArrayMarkReleased(&index_batch); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*struct_array, &index_batch)); + ScopeGuard index_guard([&index_batch]() { ArrowArrayRelease(&index_batch); }); + PAIMON_RETURN_NOT_OK(building_index_->AddBatch(&index_batch, write_batch.offset_range)); + } uint64_t memory_usage = ArrowUtils::GetArrayMemoryUsage(struct_array->data()); if (statistics) { memory_usage += ArrowUtils::GetArrayMemoryUsage(statistics->min_values->data()) + @@ -708,10 +943,12 @@ ArrowRealtimeStore::FreezeBuildingSegment() { if (building_batches_.empty()) { return std::optional>(); } - memory_segment = - std::make_shared(building_range_.value(), std::move(building_batches_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr next_index, CreateBuildingIndex()); + memory_segment = std::make_shared( + building_range_.value(), std::move(building_batches_), std::move(building_index_)); + building_index_ = std::move(next_index); sealed_segments_.push_back(memory_segment); - sealed_memory_usage_ += building_memory_usage_; + sealed_memory_usage_ += memory_segment->GetMemoryUsage(); sealed_row_count_ += building_row_count_; building_batches_.clear(); building_range_.reset(); @@ -762,9 +999,10 @@ Result> ArrowRealtimeStore::SpillFrozenSe writer.reset(); auto spill_file = std::make_shared(io_manager_, spill_file_system_, std::move(channel_id)); + // TODO(xinyu.lxy): Real-time indexes do not support spill yet. auto spilled_segment = std::make_shared( memory_segment->GetOffsetRange(), memory_segment->GetRowCount(), std::move(spilled_batches), - std::move(spill_file)); + std::move(spill_file), memory_segment->GetIndex()); cleanup_guard.Release(); { @@ -829,7 +1067,7 @@ Result> ArrowRealtimeStore::AcquireReadView() std::vector> segments = sealed_segments_; if (!building_batches_.empty()) { segments.push_back(std::make_shared( - building_range_.value(), std::vector(building_batches_))); + building_range_.value(), std::vector(building_batches_), building_index_)); } std::shared_ptr view = std::make_shared(std::move(segments)); return std::shared_ptr(std::move(view)); @@ -847,6 +1085,9 @@ Result>> ArrowRealtimeStore::CreateQuer if (context.read_batch_size <= 0) { return Status::Invalid("real-time query read batch size must be positive"); } + if (context.offset_range.begin < 0 || context.offset_range.begin > context.offset_range.end) { + return Status::Invalid("real-time query offset range is invalid"); + } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, arrow::ImportSchema(context.read_schema)); std::vector> readers; @@ -916,9 +1157,15 @@ Result>> ArrowRealtimeStore::CreateQuer return SliceQueryReaders(std::move(readers), context.read_batch_size, arrow_pool_); } + OffsetRange visible_range = arrow_view->GetOffsetRange().value(); + visible_range.begin = std::max(visible_range.begin, context.offset_range.begin); + visible_range.end = std::min(visible_range.end, context.offset_range.end); + if (visible_range.begin >= visible_range.end) { + return readers; + } std::unique_ptr reader = std::make_unique( - arrow_view, read_schema, predicate_filter, std::move(statistics_mapping), arrow_pool_, - memory_pool_); + arrow_view, read_schema, predicate_filter, std::move(statistics_mapping), context.predicate, + visible_range, arrow_pool_, memory_pool_); readers.push_back(std::move(reader)); return SliceQueryReaders(std::move(readers), context.read_batch_size, arrow_pool_); } @@ -961,13 +1208,15 @@ Result> ArrowRealtimeStore::RetireCommittedSeg RealtimeStoreDataUsage ArrowRealtimeStore::GetDataUsage() const { std::lock_guard lock(mutex_); - return RealtimeStoreDataUsage{building_memory_usage_, sealed_memory_usage_, building_row_count_, - sealed_row_count_}; + const uint64_t index_memory = building_index_ ? building_index_->GetMemoryUsage() : 0; + return RealtimeStoreDataUsage{building_memory_usage_ + index_memory, sealed_memory_usage_, + building_row_count_, sealed_row_count_}; } uint64_t ArrowRealtimeStore::GetMemoryUsage() const { std::lock_guard lock(mutex_); - return building_memory_usage_ + sealed_memory_usage_; + return building_memory_usage_ + sealed_memory_usage_ + + (building_index_ ? building_index_->GetMemoryUsage() : 0); } } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store.h b/src/paimon/core/realtime/arrow_realtime_store.h index bebaac499..d0f8dfa11 100644 --- a/src/paimon/core/realtime/arrow_realtime_store.h +++ b/src/paimon/core/realtime/arrow_realtime_store.h @@ -20,6 +20,7 @@ #pragma once #include +#include #include #include #include @@ -40,16 +41,18 @@ class FileSystem; class IOManager; class MemoryPool; class PredicateFilter; +class RealtimeIndexSet; /// Internal Arrow-backed implementation of the default `RealtimeStore`. class ArrowRealtimeStore final : public RealtimeStoreTwoPhaseSeal { public: - ArrowRealtimeStore(const std::shared_ptr& write_schema, RealtimeStoreMode mode, - StatisticsMode statistics_mode, const std::string& temp_directory, - const std::shared_ptr& spill_file_system, - const std::string& spill_compression, int32_t spill_compression_level, - const std::shared_ptr& memory_pool, - const std::shared_ptr& arrow_pool); + static Result> Create( + const std::shared_ptr& write_schema, RealtimeStoreMode mode, + StatisticsMode statistics_mode, const std::string& temp_directory, + const std::shared_ptr& spill_file_system, const std::string& spill_compression, + int32_t spill_compression_level, const std::map& index_options, + const std::shared_ptr& memory_pool, + const std::shared_ptr& arrow_pool); Status Write(RealtimeWriteBatch&& write_batch) override; @@ -79,6 +82,16 @@ class ArrowRealtimeStore final : public RealtimeStoreTwoPhaseSeal { uint64_t GetMemoryUsage() const override; private: + ArrowRealtimeStore(const std::shared_ptr& write_schema, RealtimeStoreMode mode, + StatisticsMode statistics_mode, + const std::shared_ptr& spill_file_system, + const std::string& spill_compression, int32_t spill_compression_level, + const std::map& index_options, + const std::shared_ptr& io_manager, + std::shared_ptr building_index, + const std::shared_ptr& memory_pool, + const std::shared_ptr& arrow_pool); + struct BatchStatistics { // Arrow buffers retain a raw MemoryPool pointer. Keep their allocator alive until all // statistics arrays have been released. @@ -113,6 +126,8 @@ class ArrowRealtimeStore final : public RealtimeStoreTwoPhaseSeal { class SpillBatchReader; class AppendQueryBatchReader; + Result> CreateBuildingIndex() const; + Result> CollectStatistics( const std::shared_ptr& data) const; @@ -132,7 +147,9 @@ class ArrowRealtimeStore final : public RealtimeStoreTwoPhaseSeal { std::shared_ptr spill_file_system_; std::string spill_compression_; int32_t spill_compression_level_; + std::map index_options_; mutable std::mutex mutex_; + std::shared_ptr building_index_; std::vector building_batches_; std::vector> sealed_segments_; std::optional building_range_; diff --git a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp index 07c1ba240..75e4cfd17 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_factory.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_factory.cpp @@ -20,6 +20,7 @@ #include "paimon/realtime/arrow_realtime_store_factory.h" #include +#include #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" @@ -28,7 +29,6 @@ #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/realtime/arrow_realtime_store.h" -#include "paimon/core/realtime/realtime_utils.h" #include "paimon/macros.h" namespace paimon { @@ -45,12 +45,6 @@ Result> ArrowRealtimeStoreFactory::Create( } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported_schema, arrow::ImportSchema(request.write_schema.get())); - if (request.mode != RealtimeStoreMode::APPEND_ONLY && - request.mode != RealtimeStoreMode::PRIMARY_KEY) { - return Status::Invalid("invalid real-time store mode: ", - static_cast(request.mode)); - } - PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(imported_schema)); std::shared_ptr arrow_pool = GetArrowPool(request.memory_pool); PAIMON_ASSIGN_OR_RAISE(CoreOptions options, CoreOptions::FromMap(request.options)); const CompressOptions& spill_compression = options.GetSpillCompressOptions(); @@ -65,10 +59,13 @@ Result> ArrowRealtimeStoreFactory::Create( } spill_directory = request.temp_directory; } - return std::make_shared( - imported_schema, request.mode, request.statistics_mode, spill_directory, - request.file_system, spill_compression.compress, spill_compression.zstd_level, - request.memory_pool, arrow_pool); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr store, + ArrowRealtimeStore::Create(imported_schema, request.mode, request.statistics_mode, + spill_directory, request.file_system, spill_compression.compress, + spill_compression.zstd_level, request.options, + request.memory_pool, arrow_pool)); + return std::static_pointer_cast(std::move(store)); } } // namespace paimon diff --git a/src/paimon/core/realtime/arrow_realtime_store_test.cpp b/src/paimon/core/realtime/arrow_realtime_store_test.cpp index 27dd14b9d..6ff538a9e 100644 --- a/src/paimon/core/realtime/arrow_realtime_store_test.cpp +++ b/src/paimon/core/realtime/arrow_realtime_store_test.cpp @@ -73,15 +73,16 @@ class ArrowRealtimeStoreTest : public testing::Test { }); pool_ = GetDefaultPool(); arrow_pool_ = GetArrowPool(pool_); - store_ = CreateStore(StatisticsMode::NONE); + ASSERT_OK_AND_ASSIGN(store_, CreateStore(StatisticsMode::NONE)); } - std::shared_ptr CreateStore(StatisticsMode statistics_mode) const { - return std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, statistics_mode, /*temp_directory=*/"", - /*spill_file_system=*/nullptr, - /*spill_compression=*/"zstd", - /*spill_compression_level=*/1, pool_, arrow_pool_); + Result> CreateStore(StatisticsMode statistics_mode) const { + return ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, statistics_mode, + /*temp_directory=*/"", + /*spill_file_system=*/nullptr, + /*spill_compression=*/"zstd", + /*spill_compression_level=*/1, + /*index_options=*/{}, pool_, arrow_pool_); } std::unique_ptr MakeBatch(const std::string& json) const { @@ -225,7 +226,8 @@ TEST_F(ArrowRealtimeStoreTest, TestQueryReaderClipsCommittedOffsetWithBitmap) { }); { std::unique_ptr c_schema = MakeReadSchema(read_schema); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, OffsetRange(12, 15)}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, context)); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); @@ -234,10 +236,12 @@ TEST_F(ArrowRealtimeStoreTest, TestQueryReaderClipsCommittedOffsetWithBitmap) { ASSERT_TRUE(import_result.ok()); std::shared_ptr array = std::move(import_result).ValueOrDie(); ASSERT_TRUE(array->type()->Equals(arrow::struct_(read_schema->fields()))); + ASSERT_EQ(1, array->length()); } { std::unique_ptr c_schema = MakeReadSchema(read_schema); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, OffsetRange(12, 15)}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -270,7 +274,9 @@ TEST_F(ArrowRealtimeStoreTest, TestQueryReaderClipsCommittedOffsetWithBitmap) { } std::unique_ptr c_schema = MakeReadSchema(read_schema); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -344,7 +350,9 @@ TEST_F(ArrowRealtimeStoreTest, TestSealSpillsAndKeepsPinnedMemoryViewReadable) { ASSERT_EQ(std::vector({1, 2, 3}), ReadIds(committed)); std::unique_ptr read_schema = MakeReadSchema(schema_); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr, + RealtimeQueryContext context{read_schema.get(), + /*predicate=*/nullptr, + pinned_view->GetOffsetRange().value_or(OffsetRange(0, 0)), /*read_batch_size=*/1}; ASSERT_OK_AND_ASSIGN(std::vector> query_readers, store->CreateQueryReaders(pinned_view, context)); @@ -359,7 +367,9 @@ TEST_F(ArrowRealtimeStoreTest, TestSealSpillsAndKeepsPinnedMemoryViewReadable) { ASSERT_OK_AND_ASSIGN(std::shared_ptr spilled_view, store->AcquireReadView()); std::unique_ptr spilled_read_schema = MakeReadSchema(schema_); - RealtimeQueryContext spilled_context{spilled_read_schema.get(), /*predicate=*/nullptr, + RealtimeQueryContext spilled_context{spilled_read_schema.get(), + /*predicate=*/nullptr, + spilled_view->GetOffsetRange().value_or(OffsetRange(0, 0)), /*read_batch_size=*/1}; ASSERT_OK_AND_ASSIGN(std::vector> spilled_readers, store->CreateQueryReaders(spilled_view, spilled_context)); @@ -376,9 +386,11 @@ TEST_F(ArrowRealtimeStoreTest, TestSealSpillsAndKeepsPinnedMemoryViewReadable) { TEST_F(ArrowRealtimeStoreTest, TestTwoPhaseSealDoesNotResurrectReclaimedSegment) { auto directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); - auto store = std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, directory->Str(), - directory->GetFileSystem(), "zstd", 1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, + directory->Str(), directory->GetFileSystem(), "zstd", 1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_NOK_WITH_MSG(store->SpillFrozenSegment(std::make_shared()), "requires a frozen memory segment"); ASSERT_OK_AND_ASSIGN(std::optional> empty, @@ -409,9 +421,11 @@ TEST_F(ArrowRealtimeStoreTest, TestTwoPhaseSealDoesNotResurrectReclaimedSegment) TEST_F(ArrowRealtimeStoreTest, TestRetiredSpillStorageIsReleasedByCaller) { auto directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); - auto store = std::make_shared( - schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, directory->Str(), - directory->GetFileSystem(), "zstd", 1, pool_, arrow_pool_); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr store, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, + directory->Str(), directory->GetFileSystem(), "zstd", 1, + /*index_options=*/{}, pool_, arrow_pool_)); ASSERT_OK(store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 1, "a"]])"), OffsetRange(0, 1)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, store->SealForCommit()); @@ -486,7 +500,8 @@ TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { std::unique_ptr read_schema = MakeReadSchema(schema_); std::shared_ptr predicate = PredicateBuilder::GreaterThan( /*field_index=*/1, /*field_name=*/"id", FieldType::BIGINT, Literal(int64_t{5})); - RealtimeQueryContext context{read_schema.get(), predicate}; + RealtimeQueryContext context{read_schema.get(), predicate, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -499,7 +514,9 @@ TEST_F(ArrowRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingBatch) { ASSERT_EQ(std::vector({11}), ReadIds(filtered)); std::unique_ptr unfiltered_read_schema = MakeReadSchema(schema_); - RealtimeQueryContext unfiltered_context{unfiltered_read_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext unfiltered_context{unfiltered_read_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> unfiltered_readers, store->CreateQueryReaders(view, unfiltered_context)); ASSERT_EQ(1, unfiltered_readers.size()); @@ -518,7 +535,8 @@ TEST_F(ArrowRealtimeStoreTest, TestMissingStatisticsRetainsNonMatchingBatch) { std::unique_ptr read_schema = MakeReadSchema(schema_); std::shared_ptr predicate = PredicateBuilder::GreaterThan( /*field_index=*/1, /*field_name=*/"id", FieldType::BIGINT, Literal(int64_t{5})); - RealtimeQueryContext context{read_schema.get(), predicate}; + RealtimeQueryContext context{read_schema.get(), predicate, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store_->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -528,12 +546,179 @@ TEST_F(ArrowRealtimeStoreTest, TestMissingStatisticsRetainsNonMatchingBatch) { ASSERT_EQ(std::vector({0, 1, 10, 11}), ReadIds(retained)); } +TEST_F(ArrowRealtimeStoreTest, TestHashIndexPinsHistoryAndUsesBothOffsetBounds) { + ArrowRealtimeStoreFactory factory; + RealtimeStoreCreateRequest request; + request.write_schema = MakeReadSchema(schema_); + request.memory_pool = pool_; + request.options = {{"realtime.index.hash.columns", "id,value"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, factory.Create(std::move(request))); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[10, 1, "a"], [13, 2, "a"], [20, 1, "b"]])"), OffsetRange(10, 21)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr old_view, store->AcquireReadView()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[25, 1, "a"]])"), OffsetRange(25, 26)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK( + store->Write(RealtimeWriteBatch{MakeBatch(R"([[30, 1, "a"]])"), OffsetRange(30, 31)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr current_view, store->AcquireReadView()); + ASSERT_OK(store->AdvanceCommittedOffset(26)); + + std::unique_ptr read_schema = MakeReadSchema(schema_); + const std::shared_ptr key = + PredicateBuilder::Equal(1, "id", FieldType::BIGINT, Literal(int64_t{1})); + RealtimeQueryContext old_context{read_schema.get(), key, OffsetRange(0, 100), 1024}; + ASSERT_OK_AND_ASSIGN(std::vector> old_readers, + store->CreateQueryReaders(old_view, old_context)); + ASSERT_EQ(1, old_readers.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr old_rows, + ReadResultCollector::CollectResult(std::move(old_readers[0]))); + ASSERT_EQ(std::vector({1, 1}), ReadIds(old_rows)); + + read_schema = MakeReadSchema(schema_); + const std::shared_ptr value = + PredicateBuilder::Equal(2, "value", FieldType::STRING, Literal(FieldType::STRING, "a", 1)); + RealtimeQueryContext context{read_schema.get(), value, OffsetRange(11, 30), 1}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(current_view, context)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, + ReadResultCollector::CollectResult(std::move(readers[0]))); + ASSERT_EQ(std::vector({2, 1}), ReadIds(rows)); +} + +TEST_F(ArrowRealtimeStoreTest, TestQueryResultsWithOffsetBounds) { + for (bool indexed : {false, true}) { + std::map options; + if (indexed) { + options["realtime.index.hash.columns"] = "id"; + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, "", nullptr, "zstd", + 1, options, pool_, arrow_pool_)); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch( + R"([[10, 1, "a"], [13, 2, "b"], [20, 1, "c"], [25, 2, "d"], [30, 1, "e"], [35, 1, "f"]])"), + OffsetRange(10, 36)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + const std::shared_ptr key = + PredicateBuilder::Equal(1, "id", FieldType::BIGINT, Literal(int64_t{1})); + const std::shared_ptr read_schema = + arrow::schema({schema_->field(0), schema_->field(2)}); + const std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_( + {arrow::field(SpecialFields::RealtimeOffset().Name(), arrow::int64()), + arrow::field("value", arrow::utf8())}), + indexed ? R"([[20, "c"], [30, "e"]])" + : R"([[13, "b"], [20, "c"], [25, "d"], [30, "e"]])") + .ValueOrDie(); + for (int32_t batch_size : {2, 1024}) { + std::unique_ptr c_schema = MakeReadSchema(read_schema); + RealtimeQueryContext context{c_schema.get(), key, OffsetRange(11, 31), batch_size}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, context)); + ASSERT_EQ(1, readers.size()); + readers[0] = std::make_unique( + std::move(readers[0]), context.offset_range, /*keep_offset=*/true); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, + ReadResultCollector::CollectResult(std::move(readers[0]))); + ASSERT_TRUE(rows->Equals(arrow::ChunkedArray(expected))) << rows->ToString(); + } + } +} + +TEST_F(ArrowRealtimeStoreTest, TestQueryBoundsExcludeUnreclaimedPrefixAndLaterWrites) { + for (bool indexed : {false, true}) { + std::map options; + if (indexed) { + options["realtime.index.hash.columns"] = "id"; + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::APPEND_ONLY, + StatisticsMode::NONE, "", nullptr, "zstd", + 1, options, pool_, arrow_pool_)); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[10, 1, "disk-a"], [13, 1, "disk-b"]])"), OffsetRange(10, 14)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[20, 1, "memory"]])"), OffsetRange(20, 21)})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + ASSERT_OK(store->Write( + RealtimeWriteBatch{MakeBatch(R"([[25, 1, "later"]])"), OffsetRange(25, 26)})); + // Disk covers [0, 14); the pinned view excludes offset 25, while old memory remains. + ASSERT_EQ(2, store->GetDataUsage().sealed_row_count); + ASSERT_EQ(2, store->GetDataUsage().building_row_count); + std::unique_ptr read_schema = MakeReadSchema(schema_); + const std::shared_ptr key = + PredicateBuilder::Equal(1, "id", FieldType::BIGINT, Literal(int64_t{1})); + RealtimeQueryContext context{read_schema.get(), key, OffsetRange(14, 21)}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, context)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, + ReadResultCollector::CollectResult(std::move(readers[0]))); + ASSERT_EQ(std::vector({1}), ReadIds(rows)); + const std::shared_ptr batch = + checked_pointer_cast(rows->chunk(0)); + const std::shared_ptr values = + checked_pointer_cast(batch->GetFieldByName("value")); + ASSERT_EQ("memory", values->GetString(0)); + } +} + +TEST_F(ArrowRealtimeStoreTest, TestColumnsOptionsRejectIndexesForPrimaryKeyStore) { + ASSERT_NOK_WITH_MSG( + ArrowRealtimeStore::Create(schema_, RealtimeStoreMode::PRIMARY_KEY, StatisticsMode::NONE, + "", nullptr, "zstd", 1, {{"realtime.index.hash.columns", "id"}}, + pool_, arrow_pool_), + "real-time query indexes require an append-store delegate"); +} + +TEST_F(ArrowRealtimeStoreTest, TestHashIndexRemainsInMemoryAcrossSpill) { + auto directory = UniqueTestDirectory::Create(); + ASSERT_NE(nullptr, directory); + ArrowRealtimeStoreFactory factory; + RealtimeStoreCreateRequest request; + request.write_schema = MakeReadSchema(schema_); + request.memory_pool = pool_; + request.options = {{"realtime.index.hash.columns", "id,value"}, + {Options::REALTIME_SPILL_ENABLED, "true"}}; + request.temp_directory = directory->Str(); + request.file_system = directory->GetFileSystem(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, factory.Create(std::move(request))); + ASSERT_OK(store->Write(RealtimeWriteBatch{ + MakeBatch(R"([[10, 1, "a"], [13, 2, "b"], [20, 1, "c"]])"), OffsetRange(10, 21)})); + ASSERT_OK_AND_ASSIGN(std::optional> segment, + store->SealForCommit()); + ASSERT_TRUE(segment.has_value()); + ASSERT_GT(store->GetDataUsage().sealed_memory_bytes, 0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); + ASSERT_OK(store->AdvanceCommittedOffset(21)); + std::unique_ptr read_schema = MakeReadSchema(schema_); + const std::shared_ptr key = + PredicateBuilder::Equal(1, "id", FieldType::BIGINT, Literal(int64_t{1})); + RealtimeQueryContext context{read_schema.get(), key, OffsetRange(11, 21), 1024}; + ASSERT_OK_AND_ASSIGN(std::vector> readers, + store->CreateQueryReaders(view, context)); + ASSERT_EQ(1, readers.size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr rows, + ReadResultCollector::CollectResult(std::move(readers[0]))); + ASSERT_EQ(std::vector({1}), ReadIds(rows)); +} + TEST_F(ArrowRealtimeStoreTest, TestRejectsHandlesFromAnotherStoreImplementation) { ASSERT_NOK_WITH_MSG(store_->CreateCommitReaders(std::make_shared()), "segment was not created by the Arrow real-time store"); std::unique_ptr read_schema = MakeReadSchema(schema_); - RealtimeQueryContext context{read_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{read_schema.get(), + /*predicate=*/nullptr, OffsetRange(0, 0)}; ASSERT_NOK_WITH_MSG(store_->CreateQueryReaders(std::make_shared(), context), "read view was not created by the Arrow real-time store"); diff --git a/src/paimon/core/realtime/index/hash/hash_realtime_index.cpp b/src/paimon/core/realtime/index/hash/hash_realtime_index.cpp new file mode 100644 index 000000000..325d2d7ec --- /dev/null +++ b/src/paimon/core/realtime/index/hash/hash_realtime_index.cpp @@ -0,0 +1,200 @@ +/* + * 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/index/hash/hash_realtime_index.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/predicate/literal_converter.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/core/realtime/realtime_utils.h" +#include "paimon/defs.h" +#include "paimon/macros.h" +#include "tbb/concurrent_vector.h" + +namespace paimon { + +struct HashRealtimeIndex::Posting { + tbb::concurrent_vector offsets; + std::atomic published_size{0}; +}; + +HashRealtimeIndex::Reader::Reader(std::shared_ptr index, + const OffsetRange& offset_range) + : index_(std::move(index)), offset_range_(offset_range) {} + +Result> HashRealtimeIndex::Reader::VisitEqual( + const Literal& literal) { + return VisitIn({literal}); +} + +Result> HashRealtimeIndex::Reader::VisitIn( + const std::vector& literals) { + RoaringBitmap64 result; + if (offset_range_.Empty()) { + return std::optional(std::move(result)); + } + for (const Literal& literal : literals) { + if (literal.IsNull()) { + continue; + } + if (literal.GetType() != index_->field_type_) { + return Status::Invalid("real-time hash index literal type mismatch: expected ", + FieldTypeUtils::FieldTypeToString(index_->field_type_), + ", actual ", + FieldTypeUtils::FieldTypeToString(literal.GetType())); + } + const std::optional> posting = + index_->postings_.Find(literal.HashCode()); + if (!posting) { + continue; + } + const Posting& entry = *posting.value(); + const size_t published_size = entry.published_size.load(std::memory_order_acquire); + const auto published_end = entry.offsets.cbegin() + published_size; + auto begin = std::lower_bound(entry.offsets.cbegin(), published_end, offset_range_.begin); + const auto end = std::lower_bound(begin, published_end, offset_range_.end); + // TBB posting segments are not contiguous. + for (; begin != end; ++begin) { + result.Add(*begin); + } + } + return std::optional(std::move(result)); +} + +HashRealtimeIndex::HashRealtimeIndex(std::shared_ptr schema, FieldType field_type) + : schema_(std::move(schema)), field_type_(field_type) {} + +Result> HashRealtimeIndex::Create( + const std::shared_ptr& schema, const std::shared_ptr& pool) { + if (!pool) { + return Status::Invalid("real-time hash index requires a memory pool"); + } + PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(schema)); + if (schema->num_fields() != 2 || + schema->field(0)->name() != SpecialFields::RealtimeOffset().Name() || + schema->field(1)->name() == SpecialFields::RealtimeOffset().Name()) { + return Status::Invalid( + "real-time hash index requires offset and exactly one indexed field"); + } + PAIMON_ASSIGN_OR_RAISE(FieldType type, FieldTypeUtils::ConvertToFieldType(schema->field(1))); + switch (type) { + case FieldType::BOOLEAN: + case FieldType::TINYINT: + case FieldType::SMALLINT: + case FieldType::INT: + case FieldType::BIGINT: + case FieldType::FLOAT: + case FieldType::DOUBLE: + case FieldType::STRING: + case FieldType::BINARY: + case FieldType::DATE: + break; + default: + return Status::NotImplemented("unsupported real-time hash index field type: ", + schema->field(1)->type()->ToString()); + } + return std::shared_ptr(new HashRealtimeIndex(schema, type)); +} + +Status HashRealtimeIndex::AddBatch(::ArrowArray* batch) { + if (!batch || !batch->release) { + return Status::Invalid("real-time index batch is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, + arrow::ImportArray(batch, arrow::struct_(schema_->fields()))); + const std::shared_ptr data = + checked_pointer_cast(imported); + const std::shared_ptr offsets = checked_pointer_cast( + data->GetFieldByName(SpecialFields::RealtimeOffset().Name())); + if (data->null_count() != 0 || offsets->null_count() != 0) { + return Status::Invalid("real-time index requires non-null offsets and rows"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector values, LiteralConverter::ConvertLiteralsFromArray( + *data->field(1), /*own_data=*/false)); + + for (int64_t row = 0; row < offsets->length(); ++row) { + const Literal& literal = values[row]; + if (literal.IsNull()) { + continue; + } + const uint64_t hash = literal.HashCode(); + std::optional> found = postings_.Find(hash); + std::shared_ptr posting; + if (found) { + posting = found.value(); + } else { + posting = std::make_shared(); + postings_.Insert(hash, posting); + posting_memory_usage_.fetch_add( + sizeof(Posting) + sizeof(uint64_t) + sizeof(std::shared_ptr), + std::memory_order_relaxed); + } + tbb::concurrent_vector& list = posting->offsets; + const size_t old_capacity = list.capacity(); + list.push_back(offsets->Value(row)); + // Publish initialized offsets; concurrent_vector::size() may include unfinished ones. + posting->published_size.store(list.size(), std::memory_order_release); + posting_memory_usage_.fetch_add((list.capacity() - old_capacity) * sizeof(int64_t), + std::memory_order_relaxed); + } + return Status::OK(); +} + +Result> HashRealtimeIndex::CreateReader( + const OffsetRange& offset_range) const { + if (offset_range.begin < 0 || offset_range.begin > offset_range.end) { + return Status::Invalid("real-time index query offset range is invalid"); + } + return std::shared_ptr(new Reader(shared_from_this(), offset_range)); +} + +uint64_t HashRealtimeIndex::GetMemoryUsage() const { + // Estimates payload/capacity; excludes TBB buckets and allocator overhead. + return posting_memory_usage_.load(std::memory_order_relaxed) + sizeof(*this); +} + +const char* HashRealtimeIndexFactory::Identifier() const { + return "hash-realtime"; +} + +Result> HashRealtimeIndexFactory::Create( + ::ArrowSchema* write_schema, const std::map&, + const std::shared_ptr& pool) const { + if (!write_schema || !write_schema->release) { + return Status::Invalid("real-time index schema is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, + arrow::ImportSchema(write_schema)); + return HashRealtimeIndex::Create(schema, pool); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/hash/hash_realtime_index.h b/src/paimon/core/realtime/index/hash/hash_realtime_index.h new file mode 100644 index 000000000..44efe47f3 --- /dev/null +++ b/src/paimon/core/realtime/index/hash/hash_realtime_index.h @@ -0,0 +1,92 @@ +/* + * 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 + +#include "paimon/common/utils/concurrent_hash_map.h" +#include "paimon/realtime/index/realtime_index.h" +#include "paimon/realtime/index/realtime_index_factory.h" +#include "paimon/realtime/index/realtime_index_reader.h" + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { +class MemoryPool; +enum class FieldType; + +/// Single-writer, multiple-reader hash index: TBB map -> concurrent_vector postings. +/// Atomic lengths publish initialized prefixes without copying. +/// Temporary layout; will be replaced with more lock-friendly data structures. +class HashRealtimeIndex final : public RealtimeIndex, + public std::enable_shared_from_this { + public: + static Result> Create( + const std::shared_ptr& schema, const std::shared_ptr& pool); + + Status AddBatch(::ArrowArray* batch) override; + + Result> CreateReader( + const OffsetRange& offset_range) const override; + + uint64_t GetMemoryUsage() const override; + + private: + struct Posting; + class Reader final : public RealtimeIndexReader { + public: + Reader(std::shared_ptr index, const OffsetRange& offset_range); + + Result> VisitEqual(const Literal& literal) override; + Result> VisitIn( + const std::vector& literals) override; + + private: + std::shared_ptr index_; + OffsetRange offset_range_; + }; + using PostingMap = ConcurrentHashMap>; + + HashRealtimeIndex(std::shared_ptr schema, FieldType field_type); + + std::shared_ptr schema_; + FieldType field_type_; + PostingMap postings_; + std::atomic posting_memory_usage_{0}; +}; + +class HashRealtimeIndexFactory final : public RealtimeIndexFactory { + public: + const char* Identifier() const override; + + Result> Create( + ::ArrowSchema* write_schema, const std::map& options, + const std::shared_ptr& pool) const override; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_evaluator.cpp b/src/paimon/core/realtime/index/realtime_index_evaluator.cpp new file mode 100644 index 000000000..e0c6b3806 --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_evaluator.cpp @@ -0,0 +1,120 @@ +/* + * 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/index/realtime_index_evaluator.h" + +#include + +#include "paimon/core/realtime/index/realtime_index_set.h" +#include "paimon/macros.h" +#include "paimon/predicate/compound_predicate.h" +#include "paimon/predicate/function.h" +#include "paimon/predicate/leaf_predicate.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/realtime/index/realtime_index.h" +#include "paimon/realtime/index/realtime_index_reader.h" + +namespace paimon { + +Result> RealtimeIndexEvaluator::Evaluate( + const std::shared_ptr& indexes, const std::shared_ptr& predicate, + const OffsetRange& offset_range) { + if (offset_range.begin < 0 || offset_range.begin > offset_range.end) { + return Status::Invalid("real-time index query offset range is invalid"); + } + if (offset_range.Empty()) { + return std::optional(RoaringBitmap64()); + } + if (!indexes || !predicate) { + return std::optional(); + } + const Function::Type function = predicate->GetFunction().GetType(); + if (function == Function::Type::AND || function == Function::Type::OR) { + const std::shared_ptr compound = + std::dynamic_pointer_cast(predicate); + if (!compound || compound->Children().empty()) { + return Status::Invalid("invalid or empty real-time index compound predicate"); + } + return EvaluateCompound(indexes, compound, offset_range); + } + const std::shared_ptr leaf = std::dynamic_pointer_cast(predicate); + if (!leaf) { + return std::optional(); + } + return EvaluateLeaf(indexes, leaf, offset_range); +} + +Result> RealtimeIndexEvaluator::EvaluateCompound( + const std::shared_ptr& indexes, + const std::shared_ptr& compound, const OffsetRange& offset_range) { + const Function::Type function = compound->GetFunction().GetType(); + std::optional result; + for (const std::shared_ptr& child : compound->Children()) { + PAIMON_ASSIGN_OR_RAISE(std::optional candidates, + Evaluate(indexes, child, offset_range)); + if (!candidates) { + if (function == Function::Type::OR) { + return std::optional(); + } + continue; + } + if (!result) { + result = std::move(candidates); + } else if (function == Function::Type::AND) { + *result &= *candidates; + } else { + *result |= *candidates; + } + if (function == Function::Type::AND && result->IsEmpty()) { + break; + } + } + return result; +} + +Result> RealtimeIndexEvaluator::EvaluateLeaf( + const std::shared_ptr& indexes, const std::shared_ptr& leaf, + const OffsetRange& offset_range) { + std::optional result; + for (const std::shared_ptr& index : + indexes->FindIndexes(leaf->FieldName(), leaf->GetFieldType())) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, + index->CreateReader(offset_range)); + if (!reader) { + continue; + } + PAIMON_ASSIGN_OR_RAISE( + std::optional candidates, + PredicateUtils::VisitPredicate>(leaf, reader)); + if (!candidates) { + continue; + } + if (!result) { + result = std::move(candidates); + } else { + *result &= *candidates; + } + if (result->IsEmpty()) { + break; + } + } + return result; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_evaluator.h b/src/paimon/core/realtime/index/realtime_index_evaluator.h new file mode 100644 index 000000000..f749c0dce --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_evaluator.h @@ -0,0 +1,54 @@ +/* + * 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 "paimon/realtime/offset_range.h" +#include "paimon/result.h" +#include "paimon/utils/roaring_bitmap64.h" + +namespace paimon { +class CompoundPredicate; +class LeafPredicate; +class Predicate; +class RealtimeIndexSet; + +/// Evaluates compound predicates using field index readers. +class RealtimeIndexEvaluator { + public: + RealtimeIndexEvaluator() = delete; + + static Result> Evaluate( + const std::shared_ptr& indexes, + const std::shared_ptr& predicate, const OffsetRange& offset_range); + + private: + static Result> EvaluateCompound( + const std::shared_ptr& indexes, + const std::shared_ptr& compound, const OffsetRange& offset_range); + + static Result> EvaluateLeaf( + const std::shared_ptr& indexes, + const std::shared_ptr& leaf, const OffsetRange& offset_range); +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_factory.cpp b/src/paimon/core/realtime/index/realtime_index_factory.cpp new file mode 100644 index 000000000..411928e47 --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_factory.cpp @@ -0,0 +1,48 @@ +/* + * 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/realtime/index/realtime_index_factory.h" + +#include "paimon/core/realtime/index/hash/hash_realtime_index.h" +#include "paimon/factories/factory_creator.h" +#include "paimon/realtime/index/realtime_index.h" +#include "paimon/status.h" + +namespace paimon { +const char RealtimeIndexFactory::REALTIME_INDEX_IDENTIFIER_SUFFIX[] = "-realtime"; + +Result> RealtimeIndexFactory::Get( + const std::string& identifier, ::ArrowSchema* write_schema, + const std::map& options, const std::shared_ptr& pool) { + FactoryCreator* creator = FactoryCreator::GetInstance(); + if (!creator) { + return Status::Invalid("factory creator is null pointer"); + } + auto* factory = dynamic_cast( + creator->Create(identifier + REALTIME_INDEX_IDENTIFIER_SUFFIX)); + if (!factory) { + return Status::Invalid("unknown real-time index type: ", identifier); + } + return factory->Create(write_schema, options, pool); +} + +// Keep built-in registration with the lookup entry point for static linking. +REGISTER_PAIMON_FACTORY(HashRealtimeIndexFactory); + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_reader.cpp b/src/paimon/core/realtime/index/realtime_index_reader.cpp new file mode 100644 index 000000000..c51c886a3 --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_reader.cpp @@ -0,0 +1,106 @@ +/* + * 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/realtime/index/realtime_index_reader.h" + +#include + +#include "paimon/macros.h" + +namespace paimon { + +Result> RealtimeIndexReader::VisitIsNotNull() { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitIsNull() { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitEqual(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitNotEqual(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitLessThan(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitLessOrEqual(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitGreaterThan(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitGreaterOrEqual(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitIn( + const std::vector& literals) { + RoaringBitmap64 result; + for (const Literal& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(std::optional candidates, VisitEqual(literal)); + if (!candidates) { + return std::optional(); + } + result |= *candidates; + } + return std::optional(std::move(result)); +} + +Result> RealtimeIndexReader::VisitNotIn( + const std::vector& literals) { + std::optional result; + for (const Literal& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(std::optional candidates, VisitNotEqual(literal)); + if (!candidates) { + continue; + } + if (!result) { + result = std::move(candidates); + } else { + *result &= *candidates; + } + } + return result; +} + +Result> RealtimeIndexReader::VisitStartsWith(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitEndsWith(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitContains(const Literal&) { + return std::optional(); +} + +Result> RealtimeIndexReader::VisitLike(const Literal&) { + return std::optional(); +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_set.cpp b/src/paimon/core/realtime/index/realtime_index_set.cpp new file mode 100644 index 000000000..eadd5ba3c --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_set.cpp @@ -0,0 +1,182 @@ +/* + * 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/index/realtime_index_set.h" + +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/core/realtime/realtime_utils.h" +#include "paimon/macros.h" +#include "paimon/realtime/index/realtime_index.h" +#include "paimon/realtime/index/realtime_index_factory.h" + +namespace paimon { + +const char RealtimeIndexSet::kIndexPrefix[] = "realtime.index."; +const char RealtimeIndexSet::kColumnsSuffix[] = ".columns"; + +namespace { +bool IsColumnsOption(const std::string& key) { + return StringUtils::StartsWith(key, RealtimeIndexSet::kIndexPrefix) && + StringUtils::EndsWith(key, RealtimeIndexSet::kColumnsSuffix); +} +} // namespace + +bool RealtimeIndexSet::IsEnabled(const std::map& options) { + for (const auto& option : options) { + if (IsColumnsOption(option.first)) { + return true; + } + } + return false; +} + +Result> RealtimeIndexSet::Create( + const std::shared_ptr& write_schema, + const std::map& options, const std::shared_ptr& pool) { + if (!IsEnabled(options)) { + return std::shared_ptr(); + } + PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateOffsetField(write_schema)); + const size_t prefix_length = std::strlen(kIndexPrefix); + const size_t suffix_length = std::strlen(kColumnsSuffix); + std::multimap indexes; + for (const auto& [key, value] : options) { + if (!IsColumnsOption(key)) { + continue; + } + if (key.size() <= prefix_length + suffix_length) { + return Status::Invalid("invalid real-time index option: ", key); + } + const std::string identifier = + key.substr(prefix_length, key.size() - prefix_length - suffix_length); + std::set declared; + for (std::string name : StringUtils::Split(value, ",", /*ignore_empty=*/true)) { + StringUtils::Trim(&name); + if (name.empty()) { + return Status::Invalid(key, " contains an empty field"); + } + if (!declared.insert(name).second) { + continue; + } + const int32_t position = write_schema->GetFieldIndex(name); + if (position < 0 || name == SpecialFields::RealtimeOffset().Name()) { + return Status::Invalid("invalid real-time index field: ", name); + } + const std::shared_ptr& field = write_schema->field(position); + PAIMON_ASSIGN_OR_RAISE(FieldType field_type, FieldTypeUtils::ConvertToFieldType(field)); + const std::shared_ptr index_schema = arrow::schema( + {write_schema->GetFieldByName(SpecialFields::RealtimeOffset().Name()), field}); + ArrowSchema c_schema; + ArrowSchemaMarkReleased(&c_schema); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*index_schema, &c_schema)); + ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr index, + RealtimeIndexFactory::Get(identifier, &c_schema, options, pool)); + if (!index) { + return Status::Invalid("real-time index factory returned a null index"); + } + indexes.emplace(name, Entry{position, field_type, std::move(index)}); + } + if (declared.empty()) { + return Status::Invalid(key, " must not be empty"); + } + } + return std::shared_ptr( + new RealtimeIndexSet(write_schema, std::move(indexes))); +} + +Status RealtimeIndexSet::AddBatch(::ArrowArray* batch, const OffsetRange& offset_range) { + if (!batch || !batch->release) { + return Status::Invalid("real-time index batch is null"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr imported, + arrow::ImportArray(batch, arrow::struct_(schema_->fields()))); + const std::shared_ptr data = + checked_pointer_cast(imported); + if (data->null_count() != 0) { + return Status::Invalid("real-time index requires non-null rows"); + } + const std::shared_ptr offset_field = + schema_->GetFieldByName(SpecialFields::RealtimeOffset().Name()); + const std::shared_ptr offsets = + checked_pointer_cast(data->GetFieldByName(offset_field->name())); + if (offset_range.begin < 0 || offset_range.begin >= offset_range.end || + offsets->null_count() != 0) { + return Status::Invalid("real-time index requires valid offsets"); + } + int64_t previous = offset_range.begin - 1; + for (int64_t row = 0; row < offsets->length(); ++row) { + const int64_t offset = offsets->Value(row); + if (offset <= previous || offset >= offset_range.end) { + return Status::Invalid("real-time index offsets must increase within batch bounds"); + } + previous = offset; + } + if (offset_range.begin < end_offset_) { + return Status::Invalid("real-time index batches must have ordered disjoint bounds"); + } + for (const auto& item : indexes_) { + const Entry& entry = item.second; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr projected, + arrow::StructArray::Make({offsets, data->field(entry.position)}, + {offset_field, schema_->field(entry.position)})); + ArrowArray index_batch; + ArrowArrayMarkReleased(&index_batch); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*projected, &index_batch)); + ScopeGuard batch_guard([&index_batch]() { ArrowArrayRelease(&index_batch); }); + PAIMON_RETURN_NOT_OK(entry.index->AddBatch(&index_batch)); + } + end_offset_ = offset_range.end; + return Status::OK(); +} + +std::vector> RealtimeIndexSet::FindIndexes( + const std::string& field_name, FieldType field_type) const { + std::vector> result; + const auto range = indexes_.equal_range(field_name); + for (auto iter = range.first; iter != range.second; ++iter) { + if (iter->second.type == field_type) { + result.push_back(iter->second.index); + } + } + return result; +} + +uint64_t RealtimeIndexSet::GetMemoryUsage() const { + uint64_t bytes = 0; + for (const auto& item : indexes_) { + bytes += item.second.index->GetMemoryUsage(); + } + return bytes; +} + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_set.h b/src/paimon/core/realtime/index/realtime_index_set.h new file mode 100644 index 000000000..dde885a6c --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_set.h @@ -0,0 +1,81 @@ +/* + * 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/defs.h" +#include "paimon/realtime/offset_range.h" +#include "paimon/result.h" + +struct ArrowArray; + +namespace arrow { +class Schema; +} // namespace arrow + +namespace paimon { +class MemoryPool; +class RealtimeIndex; + +/// Shared per-segment indexes with fixed field bindings. +class RealtimeIndexSet { + public: + /// Configure realtime.index..columns=field1,field2. + /// Example: realtime.index.hash.columns=id,value. + static const char kIndexPrefix[]; + static const char kColumnsSuffix[]; + + static bool IsEnabled(const std::map& options); + + static Result> Create( + const std::shared_ptr& write_schema, + const std::map& options, const std::shared_ptr& pool); + + /// Consumes a store batch and dispatches offset/field pairs to indexes. + Status AddBatch(::ArrowArray* batch, const OffsetRange& offset_range); + + std::vector> FindIndexes(const std::string& field_name, + FieldType field_type) const; + + uint64_t GetMemoryUsage() const; + + private: + struct Entry { + int32_t position; + FieldType type; + std::shared_ptr index; + }; + + RealtimeIndexSet(std::shared_ptr schema, + std::multimap&& indexes) + : schema_(std::move(schema)), indexes_(std::move(indexes)) {} + + std::shared_ptr schema_; + std::multimap indexes_; + int64_t end_offset_ = 0; +}; + +} // namespace paimon diff --git a/src/paimon/core/realtime/index/realtime_index_test.cpp b/src/paimon/core/realtime/index/realtime_index_test.cpp new file mode 100644 index 000000000..027fd7a44 --- /dev/null +++ b/src/paimon/core/realtime/index/realtime_index_test.cpp @@ -0,0 +1,375 @@ +/* + * 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/realtime/index/realtime_index.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/realtime/index/realtime_index_evaluator.h" +#include "paimon/core/realtime/index/realtime_index_set.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/realtime/index/realtime_index_factory.h" +#include "paimon/realtime/index/realtime_index_reader.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +class RealtimeIndexTest : public testing::Test { + public: + void SetUp() override { + schema_ = arrow::schema( + {DataField::ConvertDataFieldToArrowField(SpecialFields::RealtimeOffset()), + arrow::field("id", arrow::int64()), arrow::field("value", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN(index_, CreateIndex({"id", "value"})); + } + + Result> CreateIndex( + const std::vector& fields) const { + std::string columns; + for (const std::string& field : fields) { + if (!columns.empty()) { + columns.append(","); + } + columns.append(field); + } + return RealtimeIndexSet::Create(schema_, {{"realtime.index.hash.columns", columns}}, + GetDefaultPool()); + } + + Result> Evaluate(const std::shared_ptr& predicate, + const OffsetRange& range) const { + return RealtimeIndexEvaluator::Evaluate(index_, predicate, range); + } + + Status Add(const std::string& json, const OffsetRange& range) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr array, + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema_->fields()), json)); + ArrowArray batch; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*array, &batch)); + ScopeGuard batch_guard([&batch]() { ArrowArrayRelease(&batch); }); + return index_->AddBatch(&batch, range); + } + + std::shared_ptr EqualId(int64_t id) const { + // Query positions intentionally differ from write positions. + return PredicateBuilder::Equal(0, "id", FieldType::BIGINT, Literal(id)); + } + + std::shared_ptr EqualValue(const std::string& value) const { + return PredicateBuilder::Equal(0, "value", FieldType::STRING, + Literal(FieldType::STRING, value.data(), value.size())); + } + + protected: + std::shared_ptr schema_; + std::shared_ptr index_; +}; + +TEST_F(RealtimeIndexTest, RetainsAllKeyAndValueOffsetsWithinBothBounds) { + ASSERT_OK(Add(R"([[10, 1, "a"], [13, 2, "a"], [20, 1, "b"]])", OffsetRange(10, 21))); + ASSERT_OK_AND_ASSIGN(std::optional keys, + Evaluate(EqualId(1), OffsetRange(0, 21))); + ASSERT_TRUE(keys.has_value()); + ASSERT_EQ(RoaringBitmap64::From({10, 20}), *keys); + ASSERT_OK_AND_ASSIGN(std::optional values, + Evaluate(EqualValue("a"), OffsetRange(11, 20))); + ASSERT_TRUE(values.has_value()); + ASSERT_EQ(RoaringBitmap64::From({13}), *values); + + const uint64_t initial_memory = index_->GetMemoryUsage(); + ASSERT_OK(Add(R"([[21, 1, "a"], [25, 3, "a"]])", OffsetRange(21, 26))); + // Appends must not change a bounded query's result. + ASSERT_OK_AND_ASSIGN(values, Evaluate(EqualValue("a"), OffsetRange(11, 20))); + ASSERT_TRUE(values.has_value()); + ASSERT_EQ(RoaringBitmap64::From({13}), *values); + ASSERT_OK_AND_ASSIGN(keys, Evaluate(EqualId(1), OffsetRange(20, 25))); + ASSERT_TRUE(keys.has_value()); + ASSERT_EQ(RoaringBitmap64::From({20, 21}), *keys); + ASSERT_GE(index_->GetMemoryUsage(), initial_memory); +} + +TEST_F(RealtimeIndexTest, InAndCompoundPredicatesPreserveConservativeFallback) { + ASSERT_OK(Add(R"([[1, 1, "a"], [3, 2, "b"], [5, 3, null]])", OffsetRange(1, 6))); + const std::shared_ptr in = PredicateBuilder::In( + 0, "id", FieldType::BIGINT, {Literal(int64_t{1}), Literal(int64_t{3})}); + ASSERT_OK_AND_ASSIGN(std::optional candidates, + Evaluate(in, OffsetRange(1, 6))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({1, 5}), *candidates); + + const std::shared_ptr unsupported = + PredicateBuilder::GreaterThan(0, "id", FieldType::BIGINT, Literal(int64_t{0})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr conjunction, + PredicateBuilder::And({in, unsupported})); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(conjunction, OffsetRange(1, 6))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({1, 5}), *candidates); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disjunction, + PredicateBuilder::Or({in, unsupported})); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(disjunction, OffsetRange(1, 6))); + ASSERT_FALSE(candidates.has_value()); + + ASSERT_OK_AND_ASSIGN(disjunction, PredicateBuilder::Or({EqualValue("a"), EqualId(2)})); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(disjunction, OffsetRange(1, 6))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({1, 3}), *candidates); + ASSERT_OK_AND_ASSIGN(conjunction, PredicateBuilder::And({EqualValue("a"), EqualId(2)})); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(conjunction, OffsetRange(1, 6))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_TRUE(candidates->IsEmpty()); +} + +TEST_F(RealtimeIndexTest, ConcurrentReadersRetainPublishedRangesDuringPostingAndDictionaryGrowth) { + ASSERT_OK(Add(R"([[0, 1, "a"], [1, 2, "a"], [2, 1, "b"]])", OffsetRange(0, 3))); + std::atomic visible_end{3}; + std::atomic finished{false}; + std::vector readers; + ScopeGuard join_readers([&]() { + finished.store(true, std::memory_order_release); + for (std::thread& reader : readers) { + reader.join(); + } + }); + for (int32_t i = 0; i < 4; ++i) { + readers.emplace_back([&]() { + const std::shared_ptr old_key = EqualId(1); + const std::shared_ptr shared_value = EqualValue("a"); + uint64_t memory_usage = index_->GetMemoryUsage(); + do { + const int64_t end = visible_end.load(std::memory_order_acquire); + ASSERT_OK_AND_ASSIGN(std::optional candidates, + Evaluate(old_key, OffsetRange(1, 3))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({2}), *candidates); + + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(shared_value, OffsetRange(1, end))); + ASSERT_TRUE(candidates.has_value()); + RoaringBitmap64 expected; + expected.Add(1); + expected.AddRange(3, end); + ASSERT_EQ(expected, *candidates); + if (end > 3) { + ASSERT_OK_AND_ASSIGN(candidates, + Evaluate(EqualId(end - 1), OffsetRange(0, end))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({end - 1}), *candidates); + } + const uint64_t updated_memory_usage = index_->GetMemoryUsage(); + ASSERT_GE(updated_memory_usage, memory_usage); + memory_usage = updated_memory_usage; + } while (!finished.load(std::memory_order_acquire)); + }); + } + for (int64_t offset = 3; offset < 1027; offset += 2) { + // Grow both postings and the dictionary. + ASSERT_OK( + Add(fmt::format(R"([[{}, 1, "a"], [{}, {}, "a"]])", offset, offset + 1, offset + 1), + OffsetRange(offset, offset + 2))); + // Publish only after indexing succeeds. + visible_end.store(offset + 2, std::memory_order_release); + } +} + +TEST_F(RealtimeIndexTest, MissingValuesEmptyRangesAndUnknownFields) { + ASSERT_OK(Add(R"([[10, 1, null]])", OffsetRange(10, 11))); + ASSERT_OK_AND_ASSIGN(std::optional candidates, + Evaluate(EqualValue("missing"), OffsetRange(0, 100))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_TRUE(candidates->IsEmpty()); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(EqualId(1), OffsetRange(10, 10))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_TRUE(candidates->IsEmpty()); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(nullptr, OffsetRange(10, 11))); + ASSERT_FALSE(candidates.has_value()); + const std::shared_ptr unknown = + PredicateBuilder::Equal(0, "unknown", FieldType::BIGINT, Literal(int64_t{1})); + ASSERT_OK_AND_ASSIGN(candidates, Evaluate(unknown, OffsetRange(10, 11))); + ASSERT_FALSE(candidates.has_value()); + ASSERT_NOK(Evaluate(EqualId(1), OffsetRange(11, 10))); + ASSERT_NOK(Evaluate(EqualId(1), OffsetRange(-1, 10))); + ASSERT_NOK(CreateIndex({"missing"})); + ASSERT_NOK(CreateIndex({})); +} + +TEST_F(RealtimeIndexTest, FieldReadersBindIndependentRangesAndShareAppendOnlyPostings) { + ASSERT_OK(Add(R"([[10, 1, "a"], [13, 2, "b"], [20, 1, "c"]])", OffsetRange(10, 21))); + std::vector> key_indexes = + index_->FindIndexes("id", FieldType::BIGINT); + const std::vector> value_indexes = + index_->FindIndexes("value", FieldType::STRING); + ASSERT_EQ(1, key_indexes.size()); + ASSERT_EQ(1, value_indexes.size()); + std::shared_ptr key_index = std::move(key_indexes[0]); + ASSERT_NE(key_index, value_indexes[0]); + ASSERT_OK_AND_ASSIGN(std::shared_ptr first, + key_index->CreateReader(OffsetRange(10, 14))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr second, + key_index->CreateReader(OffsetRange(20, 21))); + ASSERT_OK(Add(R"([[25, 1, "d"]])", OffsetRange(25, 26))); + ASSERT_OK_AND_ASSIGN(std::optional candidates, + first->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({10}), *candidates); + ASSERT_OK_AND_ASSIGN(candidates, second->VisitIn({Literal(int64_t{1}), Literal(int64_t{2})})); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({20}), *candidates); + ASSERT_OK_AND_ASSIGN(std::shared_ptr later, + key_index->CreateReader(OffsetRange(25, 26))); + ASSERT_OK_AND_ASSIGN(candidates, later->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({25}), *candidates); + ASSERT_OK_AND_ASSIGN(candidates, first->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({10}), *candidates); + ASSERT_OK_AND_ASSIGN(candidates, first->VisitGreaterThan(Literal(int64_t{0}))); + ASSERT_FALSE(candidates.has_value()); + ASSERT_NOK(key_index->CreateReader(OffsetRange(-1, 20))); + ASSERT_NOK(key_index->CreateReader(OffsetRange(20, 10))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr empty, + key_index->CreateReader(OffsetRange(10, 10))); + ASSERT_OK_AND_ASSIGN(candidates, empty->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_TRUE(candidates->IsEmpty()); + + ASSERT_TRUE(index_->FindIndexes("missing", FieldType::BIGINT).empty()); + ASSERT_TRUE(index_->FindIndexes("id", FieldType::STRING).empty()); + // Readers retain the index after its segment releases it. + key_index.reset(); + index_.reset(); + ASSERT_OK_AND_ASSIGN(candidates, second->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({20}), *candidates); +} + +TEST_F(RealtimeIndexTest, RejectsUnknownRegisteredPlugin) { + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + ASSERT_NOK_WITH_MSG(RealtimeIndexFactory::Get("missing", &c_schema, {}, GetDefaultPool()), + "unknown real-time index type"); +} + +TEST_F(RealtimeIndexTest, ColumnsOptionsSelectPluginsAndDeduplicateDeclarations) { + const std::map unrelated = {{"file-index.hash.columns", "id"}, + {"realtime.index.hash.capacity", "16"}}; + ASSERT_FALSE(RealtimeIndexSet::IsEnabled(unrelated)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr disabled, + RealtimeIndexSet::Create(schema_, unrelated, GetDefaultPool())); + ASSERT_EQ(nullptr, disabled); + + const std::map options = { + {"realtime.index.hash.columns", ",id,, id, value,,"}}; + ASSERT_TRUE(RealtimeIndexSet::IsEnabled(options)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr indexes, + RealtimeIndexSet::Create(schema_, options, GetDefaultPool())); + ASSERT_EQ(1, indexes->FindIndexes("id", FieldType::BIGINT).size()); + ASSERT_EQ(1, indexes->FindIndexes("value", FieldType::STRING).size()); + + ASSERT_NOK_WITH_MSG(RealtimeIndexSet::Create( + schema_, {{"realtime.index.missing.columns", "id"}}, GetDefaultPool()), + "unknown real-time index type: missing"); + for (const char* key : {"realtime.index.columns", "realtime.index..columns"}) { + ASSERT_NOK_WITH_MSG(RealtimeIndexSet::Create(schema_, {{key, "id"}}, GetDefaultPool()), + "invalid real-time index option"); + } +} + +TEST_F(RealtimeIndexTest, FactoryAcceptsOnlyOneIndexedFieldAndReaderNeedsNoFieldSelection) { + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); + ASSERT_NOK_WITH_MSG(RealtimeIndexFactory::Get("hash", &c_schema, {}, GetDefaultPool()), + "exactly one indexed field"); + + const std::shared_ptr field_schema = + arrow::schema({schema_->field(0), schema_->field(1)}); + ASSERT_TRUE(arrow::ExportSchema(*field_schema, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr index, + RealtimeIndexFactory::Get("hash", &c_schema, {}, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + index->CreateReader(OffsetRange(0, 10))); + ASSERT_OK_AND_ASSIGN(std::optional candidates, + reader->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_TRUE(candidates->IsEmpty()); + ASSERT_NOK_WITH_MSG(reader->VisitEqual(Literal(FieldType::STRING, "1", 1)), + "literal type mismatch: expected BIGINT, actual STRING"); + ASSERT_NOK_WITH_MSG(reader->VisitIn({Literal(int64_t{1}), Literal(FieldType::STRING, "1", 1)}), + "literal type mismatch: expected BIGINT, actual STRING"); + + const std::shared_ptr data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(field_schema->fields()), + R"([[10, 1], [30, 1]])") + .ValueOrDie(); + ArrowArray batch; + ASSERT_TRUE(arrow::ExportArray(*data, &batch).ok()); + ScopeGuard batch_guard([&batch]() { ArrowArrayRelease(&batch); }); + ASSERT_OK(index->AddBatch(&batch)); + ASSERT_OK_AND_ASSIGN(reader, index->CreateReader(OffsetRange(11, 31))); + ASSERT_OK_AND_ASSIGN(candidates, reader->VisitEqual(Literal(int64_t{1}))); + ASSERT_TRUE(candidates.has_value()); + ASSERT_EQ(RoaringBitmap64::From({30}), *candidates); +} + +TEST_F(RealtimeIndexTest, RejectsDecimalAndTimestampFields) { + const std::vector> types = { + arrow::decimal128(10, 2), + arrow::timestamp(arrow::TimeUnit::SECOND), + arrow::timestamp(arrow::TimeUnit::MILLI), + arrow::timestamp(arrow::TimeUnit::MICRO), + arrow::timestamp(arrow::TimeUnit::NANO), + arrow::timestamp(arrow::TimeUnit::MICRO, "UTC")}; + for (const std::shared_ptr& type : types) { + const std::shared_ptr schema = + arrow::schema({schema_->field(0), arrow::field("value", type)}); + ASSERT_NOK_WITH_MSG( + RealtimeIndexSet::Create(schema, {{"realtime.index.hash.columns", "value"}}, + GetDefaultPool()), + "unsupported real-time hash index field type"); + } +} + +TEST_F(RealtimeIndexTest, FrameworkRejectsUnorderedAndOutOfBoundsOffsetsBeforePublication) { + ASSERT_NOK(Add(R"([[10, 1, "a"], [10, 2, "b"]])", OffsetRange(10, 11))); + ASSERT_NOK(Add(R"([[9, 1, "a"]])", OffsetRange(10, 11))); + ASSERT_NOK(Add(R"([[11, 1, "a"]])", OffsetRange(10, 11))); + ASSERT_NOK(Add(R"([[null, 1, "a"]])", OffsetRange(10, 11))); + ASSERT_OK(Add(R"([[10, 1, "a"]])", OffsetRange(10, 11))); + ASSERT_NOK(Add(R"([[10, 2, "b"]])", OffsetRange(10, 11))); +} + +} // namespace +} // namespace paimon::test diff --git a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp index d393f2106..6cc76a80e 100644 --- a/src/paimon/core/realtime/primary_key_realtime_store_test.cpp +++ b/src/paimon/core/realtime/primary_key_realtime_store_test.cpp @@ -72,14 +72,15 @@ std::shared_ptr NestedStoreWriteSchema() { 1)}); } -std::shared_ptr CreateStore( +Result> CreateStore( const std::shared_ptr& schema, const std::shared_ptr& pool, StatisticsMode statistics_mode = StatisticsMode::NONE, const std::string& temp_directory = "", const std::shared_ptr& spill_file_system = nullptr) { - return std::make_shared( - schema, RealtimeStoreMode::PRIMARY_KEY, statistics_mode, temp_directory, spill_file_system, - /*spill_compression=*/"zstd", - /*spill_compression_level=*/1, pool, GetArrowPool(pool)); + return ArrowRealtimeStore::Create(schema, RealtimeStoreMode::PRIMARY_KEY, statistics_mode, + temp_directory, spill_file_system, + /*spill_compression=*/"zstd", + /*spill_compression_level=*/1, + /*index_options=*/{}, pool, GetArrowPool(pool)); } std::unique_ptr MakeBatch(const std::string& json) { @@ -138,7 +139,8 @@ Result> ArrayFromJson(const std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool())); ASSERT_OK_AND_ASSIGN(std::optional> segment, store->SealForCommit()); ASSERT_FALSE(segment.has_value()); @@ -166,7 +168,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestWriteAndSealValidation) { } TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { - std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool())); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeBatch(R"([[6, 1, 1, 1, "before"], [5, 0, 0, 3, "three"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write( @@ -188,9 +191,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestCommitReaderPerStoredBatch) { TEST(PrimaryKeyRealtimeStoreTest, TestSpilledFileReturnsOneIndependentReaderPerStoredBatch) { std::unique_ptr temp_directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, temp_directory); - std::shared_ptr store = - CreateStore(StoreWriteSchema(), GetDefaultPool(), StatisticsMode::NONE, - temp_directory->Str(), temp_directory->GetFileSystem()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool(), StatisticsMode::NONE, + temp_directory->Str(), temp_directory->GetFileSystem())); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeBatch(R"([[6, 1, 1, 1, "before"], [5, 0, 0, 3, "three"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write( @@ -215,7 +218,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestSpilledFileReturnsOneIndependentReaderPerS ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr, + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0)), /*read_batch_size=*/1}; ASSERT_OK_AND_ASSIGN(std::vector> query_readers, store->CreateQueryReaders(view, context)); @@ -277,7 +282,8 @@ void AssertSlicedBatch(BatchReader* reader) { TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { std::shared_ptr schema = NestedStoreWriteSchema(); - std::shared_ptr store = CreateStore(schema, GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(schema, GetDefaultPool())); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeSlicedBatch( schema, @@ -295,14 +301,17 @@ TEST(PrimaryKeyRealtimeStoreTest, TestSlicedReadersExportZeroOffsets) { 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}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); AssertSlicedBatch(readers[0].get()); } TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { - std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool())); ASSERT_OK( store->Write(RealtimeWriteBatch{MakeBatch(R"([[0, 0, 4, 1, "one"]])"), OffsetRange(4, 5)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, @@ -339,7 +348,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + retained_view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(retained_view, context)); ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, ReadArray(std::move(readers))); @@ -351,7 +362,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestReclaimKeepsReadView) { } TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { - std::shared_ptr store = CreateStore(StoreWriteSchema(), GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool())); ASSERT_OK( store->Write(RealtimeWriteBatch{MakeBatch(R"([[1, 0, 0, 2, "two"]])"), OffsetRange(0, 1)})); ASSERT_OK_AND_ASSIGN(std::optional> segment, @@ -362,7 +374,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { ASSERT_OK_AND_ASSIGN(std::shared_ptr view, store->AcquireReadView()); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); - RealtimeQueryContext context{/*read_schema=*/c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{/*read_schema=*/c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(2, readers.size()); @@ -374,8 +388,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderPerStoredBatch) { } TEST(PrimaryKeyRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingStoredBatch) { - std::shared_ptr store = - CreateStore(StoreWriteSchema(), GetDefaultPool(), StatisticsMode::FULL); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(StoreWriteSchema(), GetDefaultPool(), StatisticsMode::FULL)); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeBatch(R"([[1, 0, 0, 1, "one"], [2, 0, 1, 2, "two"]])"), OffsetRange(0, 2)})); ASSERT_OK(store->Write(RealtimeWriteBatch{ @@ -386,7 +400,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestFullStatisticsPrunesNonMatchingStoredBatch ASSERT_TRUE(arrow::ExportSchema(*StoreWriteSchema(), c_schema.get()).ok()); std::shared_ptr predicate = PredicateBuilder::Equal( /*field_index=*/3, /*field_name=*/"id", FieldType::BIGINT, Literal(int64_t{10})); - RealtimeQueryContext context{c_schema.get(), predicate}; + RealtimeQueryContext context{c_schema.get(), predicate, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -414,7 +429,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryBatchOutlivesStoreAndReader) { auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*stored_schema, c_schema.get()).ok()); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_EQ(1, readers.size()); @@ -450,7 +467,8 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { 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 = CreateStoreWriteSchema(stored_value_fields); - std::shared_ptr store = CreateStore(stored_schema, GetDefaultPool()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr store, + CreateStore(stored_schema, GetDefaultPool())); ASSERT_OK(store->Write(RealtimeWriteBatch{ MakeSlicedBatch( stored_schema, @@ -469,7 +487,9 @@ TEST(PrimaryKeyRealtimeStoreTest, TestQueryReaderProjectsNestedFields) { CreateStoreWriteSchema(requested_value_fields); auto c_schema = std::make_unique(); ASSERT_TRUE(arrow::ExportSchema(*requested_schema, c_schema.get()).ok()); - RealtimeQueryContext context{c_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext context{c_schema.get(), + /*predicate=*/nullptr, + view->GetOffsetRange().value_or(OffsetRange(0, 0))}; ASSERT_OK_AND_ASSIGN(std::vector> readers, store->CreateQueryReaders(view, context)); ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, readers[0]->NextBatch()); diff --git a/src/paimon/core/realtime/realtime_append_only_writer.cpp b/src/paimon/core/realtime/realtime_append_only_writer.cpp index 22a81acdd..f16387d51 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.cpp +++ b/src/paimon/core/realtime/realtime_append_only_writer.cpp @@ -19,7 +19,9 @@ #include "paimon/core/realtime/realtime_append_only_writer.h" +#include #include +#include #include #include #include @@ -241,6 +243,23 @@ bool RealtimeAppendOnlyWriter::HasUnpreparedRealtimeData() const { return has_building_data_ || !sealed_segments_.empty(); } +void RealtimeAppendOnlyWriter::ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) { + std::vector> retired_segments; + { + std::lock_guard prepare_lock(prepare_mutex_); + std::lock_guard store_lock(realtime_store_mutex_); + auto retained_end = std::stable_partition( + sealed_segments_.begin(), sealed_segments_.end(), + [committed_end_offset](const std::shared_ptr& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + retired_segments.assign(std::make_move_iterator(retained_end), + std::make_move_iterator(sealed_segments_.end())); + sealed_segments_.erase(retained_end, sealed_segments_.end()); + } + // Release Arrow buffers and spill files without holding the writer locks. +} + std::shared_ptr RealtimeAppendOnlyWriter::GetMetrics() const { return file_writer_->GetMetrics(); } diff --git a/src/paimon/core/realtime/realtime_append_only_writer.h b/src/paimon/core/realtime/realtime_append_only_writer.h index 7ecfd4892..a4e356591 100644 --- a/src/paimon/core/realtime/realtime_append_only_writer.h +++ b/src/paimon/core/realtime/realtime_append_only_writer.h @@ -70,6 +70,7 @@ class RealtimeAppendOnlyWriter : public BatchWriter { Status Close() override; bool HasUnpreparedRealtimeData() const override; + void ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) override; std::shared_ptr GetMetrics() const override; diff --git a/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp b/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp index 9f7ddc9a0..67b198c5c 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp +++ b/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp @@ -50,10 +50,14 @@ TEST(RealtimeDeduplicateStateTest, TestAttachAndInstallSnapshot) { }); const std::shared_ptr pool = GetDefaultPool(); const std::shared_ptr arrow_pool = GetArrowPool(pool); - auto delegate = std::make_shared( - schema, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, - /*temp_directory=*/"", /*spill_file_system=*/nullptr, - /*spill_compression=*/"zstd", /*spill_compression_level=*/1, pool, arrow_pool); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr delegate, + ArrowRealtimeStore::Create(schema, RealtimeStoreMode::APPEND_ONLY, StatisticsMode::NONE, + /*temp_directory=*/"", + /*spill_file_system=*/nullptr, + /*spill_compression=*/"zstd", + /*spill_compression_level=*/1, + /*index_options=*/{}, pool, arrow_pool)); ASSERT_OK_AND_ASSIGN( std::shared_ptr store, ArrowDeduplicateRealtimeStore::Create(schema, {"id"}, delegate, pool, arrow_pool)); diff --git a/src/paimon/core/realtime/realtime_deduplicate_writer.cpp b/src/paimon/core/realtime/realtime_deduplicate_writer.cpp index 0dc3dd474..aac61e796 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_writer.cpp +++ b/src/paimon/core/realtime/realtime_deduplicate_writer.cpp @@ -18,6 +18,8 @@ #include "paimon/core/realtime/realtime_deduplicate_writer.h" +#include +#include #include #include @@ -313,6 +315,23 @@ bool RealtimeDeduplicateWriter::HasUnpreparedRealtimeData() const { return has_building_data_ || !sealed_segments_.empty(); } +void RealtimeDeduplicateWriter::ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) { + std::vector> retired_segments; + { + std::lock_guard prepare_lock(prepare_mutex_); + std::lock_guard store_lock(realtime_store_mutex_); + auto retained_end = std::stable_partition( + sealed_segments_.begin(), sealed_segments_.end(), + [committed_end_offset](const std::shared_ptr& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + retired_segments.assign(std::make_move_iterator(retained_end), + std::make_move_iterator(sealed_segments_.end())); + sealed_segments_.erase(retained_end, sealed_segments_.end()); + } + // Release Arrow buffers and spill files without holding the writer locks. +} + std::shared_ptr RealtimeDeduplicateWriter::GetMetrics() const { return file_writer_->GetMetrics(); } diff --git a/src/paimon/core/realtime/realtime_deduplicate_writer.h b/src/paimon/core/realtime/realtime_deduplicate_writer.h index 5dc6e65e6..b5dcec514 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_writer.h +++ b/src/paimon/core/realtime/realtime_deduplicate_writer.h @@ -69,6 +69,7 @@ class RealtimeDeduplicateWriter final : public BatchWriter { Status Sync() override; Status Close() override; bool HasUnpreparedRealtimeData() const override; + void ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) override; std::shared_ptr GetMetrics() const override; private: diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.cpp b/src/paimon/core/realtime/realtime_primary_key_writer.cpp index 7e04b67ae..13a612075 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.cpp +++ b/src/paimon/core/realtime/realtime_primary_key_writer.cpp @@ -18,6 +18,7 @@ #include "paimon/core/realtime/realtime_primary_key_writer.h" +#include #include #include #include @@ -293,6 +294,24 @@ bool RealtimePrimaryKeyWriter::HasUnpreparedRealtimeData() const { std::lock_guard lock(realtime_store_mutex_); return has_building_data_ || !sealed_segments_.empty(); } + +void RealtimePrimaryKeyWriter::ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) { + std::vector> retired_segments; + { + std::lock_guard prepare_lock(prepare_mutex_); + std::lock_guard store_lock(realtime_store_mutex_); + auto retained_end = std::stable_partition( + sealed_segments_.begin(), sealed_segments_.end(), + [committed_end_offset](const std::shared_ptr& segment) { + return segment->GetOffsetRange().end > committed_end_offset; + }); + retired_segments.assign(std::make_move_iterator(retained_end), + std::make_move_iterator(sealed_segments_.end())); + sealed_segments_.erase(retained_end, sealed_segments_.end()); + } + // Release Arrow buffers and spill files without holding the writer locks. +} + std::shared_ptr RealtimePrimaryKeyWriter::GetMetrics() const { return merge_tree_writer_->GetMetrics(); } diff --git a/src/paimon/core/realtime/realtime_primary_key_writer.h b/src/paimon/core/realtime/realtime_primary_key_writer.h index 56f51bd50..3a4837222 100644 --- a/src/paimon/core/realtime/realtime_primary_key_writer.h +++ b/src/paimon/core/realtime/realtime_primary_key_writer.h @@ -67,6 +67,7 @@ class RealtimePrimaryKeyWriter final : public BatchWriter { Status Sync() override; Status Close() override; bool HasUnpreparedRealtimeData() const override; + void ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) override; std::shared_ptr GetMetrics() const override; private: diff --git a/src/paimon/core/table/source/append_only_table_read.cpp b/src/paimon/core/table/source/append_only_table_read.cpp index ca786e206..86518f757 100644 --- a/src/paimon/core/table/source/append_only_table_read.cpp +++ b/src/paimon/core/table/source/append_only_table_read.cpp @@ -252,8 +252,10 @@ Result> AppendOnlyTableRead::CreateRealtimeReader( PAIMON_ASSIGN_OR_RAISE(std::shared_ptr realtime_predicate, PredicateUtils::CreatePickedFieldFilter(context_->GetPredicate(), realtime_field_name_to_index)); - RealtimeQueryContext query_context{c_read_schema.get(), std::move(realtime_predicate), - context_->GetCoreOptions().GetReadBatchSize()}; + RealtimeQueryContext query_context{ + c_read_schema.get(), std::move(realtime_predicate), + OffsetRange(realtime_split->CommittedEndOffset(), realtime_split->MemoryEndOffset()), + context_->GetCoreOptions().GetReadBatchSize()}; PAIMON_ASSIGN_OR_RAISE(std::vector> memory_readers, memory.store->CreateQueryReaders(memory.read_view, query_context)); const size_t first_memory_reader = readers.size(); diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index eeed68e59..40b415583 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -120,8 +120,10 @@ Result>> CreateMemoryReaders( PAIMON_ASSIGN_OR_RAISE(std::shared_ptr primary_key_predicate, PredicateUtils::CreatePickedFieldFilter(context->GetPredicate(), primary_key_name_to_index)); - RealtimeQueryContext query_context{c_schema.get(), std::move(primary_key_predicate), - context->GetCoreOptions().GetReadBatchSize()}; + RealtimeQueryContext query_context{ + c_schema.get(), std::move(primary_key_predicate), + OffsetRange(split->CommittedEndOffset(), split->MemoryEndOffset()), + context->GetCoreOptions().GetReadBatchSize()}; PAIMON_ASSIGN_OR_RAISE(std::vector> batch_readers, memory.store->CreateQueryReaders(memory.read_view, query_context)); PAIMON_ASSIGN_OR_RAISE( diff --git a/src/paimon/core/table/source/realtime_table_scan.cpp b/src/paimon/core/table/source/realtime_table_scan.cpp index 42a7171b3..c53776302 100644 --- a/src/paimon/core/table/source/realtime_table_scan.cpp +++ b/src/paimon/core/table/source/realtime_table_scan.cpp @@ -166,6 +166,7 @@ Result>> RealtimeTableScan::CreateRealtimeSpl PAIMON_ASSIGN_OR_RAISE(std::string opaque_ticket, realtime_context_->PinReadView(memory, read_view_ttl_millis_)); pinned_tickets.push_back(opaque_ticket); + // Use the disk snapshot's committed end; memory reclamation may lag. return std::make_shared( RealtimeSplit::kCurrentVersion, snapshot_id, key.partition, key.bucket, std::move(grouped_disk_splits), GetCommittedEndOffset(committed_offsets, key), diff --git a/src/paimon/core/utils/batch_writer.h b/src/paimon/core/utils/batch_writer.h index abe4b4c4a..b9e506290 100644 --- a/src/paimon/core/utils/batch_writer.h +++ b/src/paimon/core/utils/batch_writer.h @@ -75,11 +75,14 @@ class BatchWriter { /// Close this writer, the call will delete newly generated but not committed files. virtual Status Close() = 0; - /// Whether this writer owns real-time data not covered by a successful prepare. + /// Whether this writer owns real-time data not covered by a successful prepare or refresh. virtual bool HasUnpreparedRealtimeData() const { return false; } + /// Releases pending real-time segments fully covered by the installed committed offset. + virtual void ReleaseCommittedRealtimeSegments(int64_t committed_end_offset) {} + virtual std::shared_ptr GetMetrics() const = 0; }; diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 84741f657..c49c13a3d 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp @@ -50,6 +50,7 @@ #include "paimon/common/predicate/equal.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" @@ -79,6 +80,7 @@ #include "paimon/predicate/leaf_predicate.h" #include "paimon/predicate/predicate.h" #include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/predicate_utils.h" #include "paimon/read_context.h" #include "paimon/reader/count_reader.h" #include "paimon/realtime/arrow_realtime_store_factory.h" @@ -1110,7 +1112,10 @@ class RealtimeWriteInteTest : public ::testing::Test { PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportSchema(*schema_layout->StoreWriteSchema(), read_schema.get())); ScopeGuard schema_guard([schema = read_schema.get()]() { ArrowSchemaRelease(schema); }); - RealtimeQueryContext query_context{read_schema.get(), /*predicate=*/nullptr}; + RealtimeQueryContext query_context{ + read_schema.get(), + /*predicate=*/nullptr, + views[0].read_view->GetOffsetRange().value_or(OffsetRange(0, 0))}; PAIMON_ASSIGN_OR_RAISE( std::vector> readers, views[0].store->CreateQueryReaders(views[0].read_view, query_context)); @@ -6239,6 +6244,402 @@ TEST_F(RealtimeWriteInteTest, TestCountReaderDoesNotSupportRealtimeSplit) { ASSERT_OK(PrepareAndClose(writer.get())); } +// TODO(xinyu.lxy): Refactor the real-time hash index tests. +class RealtimeHashIndexInteTest : public RealtimeWriteInteTest { + protected: + void SetUp() override { + RealtimeWriteInteTest::SetUp(); + ASSERT_FALSE(HasFatalFailure()); + options_["realtime.index.hash.columns"] = "id,payload"; + options_[Options::REALTIME_STORE_STATS_MODE] = "none"; + options_[Options::READ_BATCH_SIZE] = "2"; + } + + std::shared_ptr EqualId(int64_t id) const { + return PredicateBuilder::Equal(schema_->GetFieldIndex("id"), "id", FieldType::BIGINT, + Literal(id)); + } + + std::shared_ptr EqualPayload(const std::string& payload) const { + return PredicateBuilder::Equal(schema_->GetFieldIndex("payload"), "payload", + FieldType::STRING, + Literal(FieldType::STRING, payload.data(), payload.size())); + } + + void CheckPlanRows(const std::shared_ptr& plan, + const std::shared_ptr& realtime_context, + const std::shared_ptr& predicate, + const std::vector& expected) const { + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_predicate, + PredicateUtils::CreatePickedFieldFilter( + predicate, {{"id", 0}, {"payload", 1}, {"pt", 2}})); + if (predicate) { + ASSERT_NE(nullptr, read_predicate); + } + ASSERT_OK_AND_ASSIGN(std::vector actual, + ReadRows(plan, realtime_context, read_predicate, + /*enable_predicate_filter=*/false)); + ASSERT_EQ(std::multiset(expected.begin(), expected.end()), + std::multiset(actual.begin(), actual.end())); + } + + void CheckQueries(const std::shared_ptr& realtime_context, + const std::vector>& predicates, + const std::vector>& expected) const { + ASSERT_EQ(predicates.size(), expected.size()); + for (size_t i = 0; i < predicates.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, + CreatePlan(realtime_context, predicates[i])); + ASSERT_NO_FATAL_FAILURE( + CheckPlanRows(plan, realtime_context, predicates[i], expected[i])); + } + } + + void CheckProjectedQuery(const std::shared_ptr& realtime_context, + const std::shared_ptr& predicate, + const std::vector& read_fields, + bool enable_predicate_filter, const std::string& expected_json) const { + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + std::map field_positions; + for (int32_t i = 0; i < static_cast(read_fields.size()); ++i) { + field_positions.emplace(read_fields[i], i); + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_predicate, + PredicateUtils::CreatePickedFieldFilter(predicate, field_positions)); + ASSERT_NE(nullptr, read_predicate); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr actual, + ReadPlan(plan, realtime_context, read_fields, read_predicate, enable_predicate_filter)); + arrow::FieldVector fields = {arrow::field("_VALUE_KIND", arrow::int8())}; + for (const std::string& name : read_fields) { + fields.push_back(schema_->GetFieldByName(name)); + } + arrow::Result> expected_result = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), expected_json); + ASSERT_TRUE(expected_result.ok()) << expected_result.status().ToString(); + const std::shared_ptr expected = std::move(expected_result).ValueOrDie(); + ASSERT_NE(nullptr, actual); + ASSERT_TRUE(std::make_shared(expected)->Equals(*actual)) + << actual->ToString(); + } +}; + +TEST_F(RealtimeHashIndexInteTest, TestKeyAndValueIndexesPruneWithoutStatisticsOrResidualFilter) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + const std::vector rows = {{1, "red", "p0"}, + {2, "blue", "p0"}, + {3, "red", "p0"}, + {4, "green", "p0"}, + {5, "blue", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeBatch(rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + const std::vector> predicates = { + EqualId(3), + PredicateBuilder::In(0, "id", FieldType::BIGINT, + {Literal(int64_t{1}), Literal(int64_t{5})}), + EqualId(99), + EqualPayload("red"), + PredicateBuilder::In( + 1, "payload", FieldType::STRING, + {Literal(FieldType::STRING, "red", 3), Literal(FieldType::STRING, "green", 5)}), + EqualPayload("missing")}; + const std::vector> expected = { + {rows[2]}, {rows[0], rows[4]}, {}, {rows[0], rows[2]}, {rows[0], rows[2], rows[3]}, {}}; + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + // Query the same postings after segment rotation, before any data file is committed. + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeHashIndexInteTest, TestIndexedFieldCombinesWithStatisticsOnAnotherField) { + fields_.push_back(arrow::field("score", arrow::int64())); + schema_ = arrow::schema(fields_); + options_["realtime.index.hash.columns"] = "id"; + options_[Options::REALTIME_STORE_STATS_MODE] = "full"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + const std::vector batches = { + // The index finds id 7, but max(score) excludes the entire batch. + R"([[0, 7, "below", "p0", 10], [1, 8, "below-other", "p0", 20]])", + // Both rows pass score statistics; only the index can remove id 8. + R"([[2, 7, "match", "p0", 150], [3, 8, "other", "p0", 160]])", + // id 7 lies inside min/max, but the index has no candidate in this batch. + R"([[4, 6, "no-match", "p0", 170], [5, 8, "no-match", "p0", 180]])"}; + for (const std::string& json : batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + MakeUnpartitionedBatchFromJson(json)); + ASSERT_OK(writer->Write(std::move(batch))); + } + const std::shared_ptr indexed = EqualId(7); + const std::shared_ptr statistics = PredicateBuilder::GreaterThan( + schema_->GetFieldIndex("score"), "score", FieldType::BIGINT, Literal(int64_t{100})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr combined, + PredicateBuilder::And({indexed, statistics})); + const std::vector> predicates = {indexed, statistics, combined}; + const std::vector expected = { + R"([[0, 7, 10], [0, 7, 150]])", R"([[0, 7, 150], [0, 8, 160], [0, 6, 170], [0, 8, 180]])", + R"([[0, 7, 150]])"}; + const auto check_candidates = [&]() { + for (size_t i = 0; i < predicates.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_NO_FATAL_FAILURE( + CheckProjectedQuery(realtime_context, predicates[i], {"id", "score"}, + /*enable_predicate_filter=*/false, expected[i])); + } + }; + ASSERT_NO_FATAL_FAILURE(check_candidates()); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_NO_FATAL_FAILURE(check_candidates()); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeHashIndexInteTest, TestDeduplicateRebuildsFileIndexesAfterSkippingDeletedRows) { + ConfigureRealtimeDeduplicateTable(); + options_["file-index.bitmap.columns"] = SpecialFields::RealtimeOffset().Name() + ",id,payload"; + options_[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = "1MB"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr batch, + MakeBatch({{1, "old", "p0"}, + {2, "deleted", "p0"}, + {1, "new", "p0"}, + {3, "keep", "p0"}, + {2, "ignored", "p0"}, + {4, "tail", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::INSERT, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::UPDATE_AFTER, RecordBatch::RowKind::INSERT, + RecordBatch::RowKind::DELETE, RecordBatch::RowKind::INSERT})); + ASSERT_OK(writer->Write(std::move(batch))); + const std::vector> predicates = { + EqualId(1), EqualId(2), + EqualId(3), EqualId(4), + EqualPayload("old"), EqualPayload("deleted"), + EqualPayload("new"), EqualPayload("keep"), + EqualPayload("tail"), nullptr}; + const std::vector> expected = { + {{1, "new", "p0"}}, + {}, + {{3, "keep", "p0"}}, + {{4, "tail", "p0"}}, + {}, + {}, + {{1, "new", "p0"}}, + {{3, "keep", "p0"}}, + {{4, "tail", "p0"}}, + {{1, "new", "p0"}, {3, "keep", "p0"}, {4, "tail", "p0"}}}; + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 6), progress[0].offset_range); + const std::vector> files = NewFiles(progress); + ASSERT_EQ(1, files.size()); + // Surviving offsets 2, 3, 5 become file row positions 0, 1, 2. + ASSERT_EQ(3, files[0]->row_count); + ASSERT_NE(nullptr, files[0]->embedded_index); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, expected)); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + // No memory store or residual filter can repair an incorrect file-index row position. + ASSERT_NO_FATAL_FAILURE(CheckQueries(/*realtime_context=*/nullptr, predicates, expected)); +} + +TEST_F(RealtimeHashIndexInteTest, TestSnapshotBoundsExcludeCommittedPrefixAndLaterPostings) { + options_["file-index.bitmap.columns"] = "payload"; + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr disk_batch, + MakeBatch({{1, "match", "p0"}, {2, "miss", "p0"}, {3, "match", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(disk_batch))); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr memory_batch, + MakeBatch({{4, "match", "p0"}, {5, "miss", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(memory_batch))); + const std::shared_ptr predicate = EqualPayload("match"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned, CreatePlan(realtime_context, predicate)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr pinned_across_refresh, + CreatePlan(realtime_context, predicate)); + ASSERT_EQ(1, pinned->Splits().size()); + const std::shared_ptr split = + std::dynamic_pointer_cast(pinned->Splits()[0]); + ASSERT_NE(nullptr, split); + ASSERT_EQ(3, split->CommittedEndOffset()); + ASSERT_EQ(5, split->MemoryEndOffset()); + { + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, + RealtimeContextImpl::Cast(realtime_context)); + ASSERT_OK_AND_ASSIGN(RealtimeReadState state, context->AcquireReadState()); + ASSERT_EQ(1, state.views.size()); + ASSERT_EQ(3, state.views[0].store->GetDataUsage().sealed_row_count); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr later_batch, + MakeBatch({{6, "match", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(later_batch))); + const std::vector old_expected = { + {1, "match", "p0"}, {3, "match", "p0"}, {4, "match", "p0"}}; + ASSERT_NO_FATAL_FAILURE(CheckPlanRows(pinned, realtime_context, predicate, old_expected)); + const std::vector> latest_expected = { + {{1, "match", "p0"}, {3, "match", "p0"}, {4, "match", "p0"}, {6, "match", "p0"}}}; + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, {predicate}, latest_expected)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_NO_FATAL_FAILURE( + CheckPlanRows(pinned_across_refresh, realtime_context, predicate, old_expected)); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, {predicate}, latest_expected)); + ASSERT_OK(PrepareAndClose(writer.get())); +} + +TEST_F(RealtimeHashIndexInteTest, TestPinnedDeduplicatePredicatesSurviveUpdateDeleteAndRefresh) { + ConfigureRealtimeDeduplicateTable(); + options_["file-index.bitmap.columns"] = SpecialFields::RealtimeOffset().Name() + ",id,payload"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK(WriteCommitAndRefresh({{1, "disk-old", "p0"}, {2, "keep", "p0"}}, + /*partitioned=*/false, /*bucket=*/0, /*row_kinds=*/{}, + /*commit_identifier=*/0, writer.get())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_update, + MakeBatch({{1, "old", "p0"}, {3, "three", "p0"}}, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(first_update))); + const std::vector> predicates = { + EqualId(1), EqualPayload("old"), EqualPayload("new"), EqualPayload("disk-old")}; + std::vector> old_plans; + for (const std::shared_ptr& predicate : predicates) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + old_plans.push_back(std::move(plan)); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_update, + MakeBatch({{1, "new", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::UPDATE_AFTER})); + ASSERT_OK(writer->Write(std::move(second_update))); + const std::vector> updated_expected = { + {{1, "new", "p0"}}, {}, {{1, "new", "p0"}}, {}}; + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, updated_expected)); + std::vector> updated_plans; + for (const std::shared_ptr& predicate : predicates) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr plan, CreatePlan(realtime_context, predicate)); + updated_plans.push_back(std::move(plan)); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr deletion, + MakeBatch({{1, "ignored", "p0"}}, /*partitioned=*/false, /*bucket=*/0, + {RecordBatch::RowKind::DELETE})); + ASSERT_OK(writer->Write(std::move(deletion))); + const std::vector> deleted_expected(predicates.size()); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, deleted_expected)); + ASSERT_OK_AND_ASSIGN(std::vector progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + const std::vector> old_expected = { + {{1, "old", "p0"}}, {{1, "old", "p0"}}, {}, {}}; + for (size_t i = 0; i < predicates.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_NO_FATAL_FAILURE( + CheckPlanRows(old_plans[i], realtime_context, predicates[i], old_expected[i])); + ASSERT_NO_FATAL_FAILURE( + CheckPlanRows(updated_plans[i], realtime_context, predicates[i], updated_expected[i])); + } + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, predicates, deleted_expected)); + ASSERT_NO_FATAL_FAILURE(CheckQueries(realtime_context, {EqualId(2), EqualId(3)}, + {{{2, "keep", "p0"}}, {{3, "three", "p0"}}})); + ASSERT_OK(writer->Close()); +} + +TEST_F(RealtimeHashIndexInteTest, TestResidualFiltersComposeWithHashCandidatesAndProjection) { + CreateTable(/*partition_keys=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, MakeUnpartitionedBatchFromJson(R"([ + [0, 1, "hit", "p0"], [1, 2, "hit", "p1"], + [2, 3, "miss", "p0"], [3, 4, null, "p1"], + [4, 5, "hit", "p0"], [5, 6, null, "p0"] + ])")); + ASSERT_OK(writer->Write(std::move(batch))); + const std::shared_ptr hit = EqualPayload("hit"); + const std::shared_ptr range = + PredicateBuilder::GreaterThan(0, "id", FieldType::BIGINT, Literal(int64_t{1})); + const std::shared_ptr unindexed = + PredicateBuilder::Equal(2, "pt", FieldType::STRING, Literal(FieldType::STRING, "p1", 2)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_range, PredicateBuilder::And({hit, range})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_range, PredicateBuilder::Or({hit, range})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr and_unindexed, + PredicateBuilder::And({hit, unindexed})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr or_unindexed, + PredicateBuilder::Or({hit, unindexed})); + const std::vector> predicates = { + and_range, + or_range, + and_unindexed, + or_unindexed, + PredicateBuilder::IsNull(1, "payload", FieldType::STRING), + PredicateBuilder::IsNotNull(1, "payload", FieldType::STRING)}; + const std::string all_candidates = R"([ + [0, "hit", 1, "p0"], [0, "hit", 2, "p1"], [0, "miss", 3, "p0"], + [0, null, 4, "p1"], [0, "hit", 5, "p0"], [0, null, 6, "p0"] + ])"; + const std::string hit_candidates = + R"([[0, "hit", 1, "p0"], [0, "hit", 2, "p1"], [0, "hit", 5, "p0"]])"; + const std::vector expected_filtered = { + R"([[0, "hit", 2, "p1"], [0, "hit", 5, "p0"]])", + all_candidates, + R"([[0, "hit", 2, "p1"]])", + R"([[0, "hit", 1, "p0"], [0, "hit", 2, "p1"], [0, null, 4, "p1"], [0, "hit", 5, "p0"]])", + R"([[0, null, 4, "p1"], [0, null, 6, "p0"]])", + R"([[0, "hit", 1, "p0"], [0, "hit", 2, "p1"], [0, "miss", 3, "p0"], [0, "hit", 5, "p0"]])"}; + for (size_t i = 0; i < predicates.size(); ++i) { + SCOPED_TRACE(i); + ASSERT_NO_FATAL_FAILURE(CheckProjectedQuery( + realtime_context, predicates[i], {"payload", "id", "pt"}, + /*enable_predicate_filter=*/false, i == 0 || i == 2 ? hit_candidates : all_candidates)); + ASSERT_NO_FATAL_FAILURE( + CheckProjectedQuery(realtime_context, predicates[i], {"payload", "id", "pt"}, + /*enable_predicate_filter=*/true, expected_filtered[i])); + } + ASSERT_NO_FATAL_FAILURE(CheckProjectedQuery(realtime_context, and_range, {"payload", "id"}, + /*enable_predicate_filter=*/true, + R"([[0, "hit", 2], [0, "hit", 5]])")); + ASSERT_NO_FATAL_FAILURE(CheckProjectedQuery(realtime_context, hit, {"payload"}, + /*enable_predicate_filter=*/false, + R"([[0, "hit"], [0, "hit"], [0, "hit"]])")); + ASSERT_NO_FATAL_FAILURE(CheckProjectedQuery(realtime_context, predicates[4], {"payload"}, + /*enable_predicate_filter=*/true, + R"([[0, null], [0, null]])")); + ASSERT_OK(PrepareAndClose(writer.get())); +} + TEST_F(RealtimeWriteInteTest, TestProjectionAndPredicateForMemoryAndDisk) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, @@ -6889,6 +7290,98 @@ TEST_F(RealtimeWriteInteTest, TestRefreshCommittedSnapshotReclaimsMemory) { ASSERT_OK(PrepareAndClose(writer.get())); } +TEST_F(RealtimeWriteInteTest, TestRefreshReleasesMemoryHeldByUnpreparedSealedSegments) { + CreateTable(/*partition_keys=*/{}); + // Keep the Arrow allocator alive until both the writer and its context have released any + // input buffers, including when an assertion returns early. + std::shared_ptr input_pool = GetMemoryPool(); + std::shared_ptr arrow_pool = GetArrowPool(input_pool); + ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr writer, + CreateRealtimeWriter(realtime_context)); + + std::vector committed_rows = MakeRows(/*first_id=*/0, /*count=*/4, /*partition=*/"p0"); + for (Row& row : committed_rows) { + std::get<1>(row) = std::string(64 * 1024, 'x'); + } + // Track the input buffers separately from store gauges, which stop counting a segment as + // soon as the store retires it even if the writer still owns its pending commit handle. + { + arrow::Int64Builder offsets(arrow_pool.get()); + arrow::Int64Builder ids(arrow_pool.get()); + arrow::StringBuilder payloads(arrow_pool.get()); + arrow::StringBuilder partitions(arrow_pool.get()); + for (const Row& row : committed_rows) { + const auto& [id, payload, partition] = row; + ASSERT_TRUE(offsets.Append(id).ok()); + ASSERT_TRUE(ids.Append(id).ok()); + ASSERT_TRUE(payloads.Append(payload).ok()); + ASSERT_TRUE(partitions.Append(partition).ok()); + } + arrow::ArrayVector columns; + const std::vector builders = {&offsets, &ids, &payloads, &partitions}; + for (arrow::ArrayBuilder* builder : builders) { + std::shared_ptr column; + ASSERT_TRUE(builder->Finish(&column).ok()); + columns.push_back(std::move(column)); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr schema_layout, + RealtimeSchemaLayout::Create(RealtimeStoreMode::APPEND_ONLY, schema_)); + auto data = + arrow::StructArray::Make(std::move(columns), schema_layout->InputSchema()->fields()); + ASSERT_TRUE(data.ok()); + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*data.ValueOrDie(), &c_array).ok()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + RecordBatchBuilder(&c_array).SetBucket(0).Finish()); + ASSERT_OK(writer->Write(std::move(batch))); + } + ASSERT_OK(writer->Seal()); + + ResetExternalOffset(/*partition=*/{}, /*bucket=*/0, /*next_offset=*/4); + std::vector tail_rows = MakeRows(/*first_id=*/4, /*count=*/2, /*partition=*/"p0"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr tail_batch, + MakeBatch(tail_rows, /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(tail_batch))); + ASSERT_OK(writer->Seal()); + + // An independent context persists the same input prefix while this writer has only sealed + // it. Do not prepare this writer before refresh: that would already release its handles. + ASSERT_OK_AND_ASSIGN(std::shared_ptr committing_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committing_writer, + CreateRealtimeWriter(committing_context)); + ResetExternalOffset(/*partition=*/{}, /*bucket=*/0, /*next_offset=*/0); + ASSERT_OK_AND_ASSIGN(std::unique_ptr committed_batch, + MakeBatch(committed_rows, /*partitioned=*/false)); + ASSERT_OK(committing_writer->Write(std::move(committed_batch))); + ASSERT_OK_AND_ASSIGN(std::vector commits, + committing_writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(commits, /*commit_identifier=*/0)); + ASSERT_OK(committing_writer->Close()); + + const uint64_t allocated_before_refresh = input_pool->CurrentUsage(); + ASSERT_GT(allocated_before_refresh, 0); + ASSERT_OK_AND_ASSIGN(uint64_t memory_before_refresh, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_LT(input_pool->CurrentUsage(), allocated_before_refresh); + ASSERT_EQ(0, input_pool->CurrentUsage()); + ASSERT_OK_AND_ASSIGN(uint64_t memory_after_refresh, GetRealtimeMemoryUsage(realtime_context)); + ASSERT_GT(memory_after_refresh, 0); + ASSERT_LT(memory_after_refresh, memory_before_refresh); + + std::vector expected_rows = committed_rows; + expected_rows.insert(expected_rows.end(), tail_rows.begin(), tail_rows.end()); + ASSERT_OK_AND_ASSIGN(std::vector rows, ReadRows(realtime_context)); + ASSERT_EQ(expected_rows, rows); + ASSERT_OK_AND_ASSIGN(std::vector remaining, + writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, remaining.size()); + ASSERT_EQ(OffsetRange(4, 6), remaining[0].offset_range); + ASSERT_OK(writer->Close()); +} + TEST_F(RealtimeWriteInteTest, TestPlanPinsMemoryAcrossRefresh) { CreateTable(/*partition_keys=*/{}); ASSERT_OK_AND_ASSIGN(std::shared_ptr realtime_context,