feat(realtime): add hash-based offset indexes with data file rolling (#444)
diff --git a/LICENSE b/LICENSE index b29f200..fc8d265 100644 --- a/LICENSE +++ b/LICENSE
@@ -207,6 +207,8 @@ * include/paimon/factories/singleton.h * src/paimon/common/factories/singleton.cpp +* src/paimon/common/utils/hash_map.h +* src/paimon/common/utils/dense_hash_table.h Copyright: 2014-present Alibaba Inc. Home page: https://havenask.net/
diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 0cb9ccb..a8a4afe 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt
@@ -405,6 +405,8 @@ core/realtime/arrow_realtime_store.cpp core/realtime/arrow_deduplicate_realtime_store.cpp 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/realtime_key_offset_index.cpp core/realtime/realtime_offset_file_index_lookup.cpp core/realtime/realtime_deduplicate_state.cpp @@ -705,6 +707,9 @@ common/utils/arrow/mem_utils_test.cpp common/utils/arrow/status_utils_test.cpp common/utils/concurrent_hash_map_test.cpp + common/utils/hash_map_test.cpp + common/utils/dense_hash_table_test.cpp + common/utils/dense_hash_table_file_reader_test.cpp common/utils/projected_row_test.cpp common/utils/projected_array_test.cpp common/utils/bit_set_test.cpp @@ -873,6 +878,7 @@ core/memory/writer_memory_manager_test.cpp 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/primary_key_realtime_store_test.cpp core/realtime/realtime_deduplicate_state_test.cpp core/realtime/realtime_key_offset_index_test.cpp
diff --git a/src/paimon/common/utils/dense_hash_table.h b/src/paimon/common/utils/dense_hash_table.h new file mode 100644 index 0000000..097978a --- /dev/null +++ b/src/paimon/common/utils/dense_hash_table.h
@@ -0,0 +1,255 @@ +/* + * Copyright 2014-present Alibaba Inc. + * + * Licensed 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. + */ + +// Adapted from Alibaba Havenask +// https://github.com/alibaba/havenask/blob/main/aios/storage/indexlib/index/common/hash_table/DenseHashTable.h + +#pragma once + +#include <algorithm> +#include <cstddef> +#include <cstdint> +#include <cstring> +#include <limits> +#include <memory> +#include <optional> +#include <type_traits> +#include <utility> + +#include "paimon/common/memory/bytes_utils.h" +#include "paimon/fs/file_system.h" +#include "paimon/macros.h" +#include "paimon/result.h" + +namespace paimon { + +template <typename Key, typename Value> +class DenseHashTableFileReader; + +/// A fixed-capacity, linear-probing hash table for already-hashed unsigned keys. +/// All key and value bit patterns are supported. +/// Mutation requires external synchronization; mounted hash tables are read-only. +/// +/// The native layout is [uint64 bucket_count, uint64 key_count], followed by fixed-size +/// [Key, Value] buckets and two special-key buckets. Integers use host byte order and fields have +/// no padding. Dumped buffers require matching key/value types and endianness when read. This is +/// a container layout, not a versioned Paimon index format. +template <typename Key, typename Value> +class DenseHashTable { + static_assert(std::is_same_v<Key, uint32_t> || std::is_same_v<Key, uint64_t>); + static_assert(std::is_integral_v<Value> && !std::is_same_v<Value, bool>); + + public: + static Result<std::unique_ptr<DenseHashTable>> Create(uint64_t capacity, int32_t occupancy_pct, + const std::shared_ptr<MemoryPool>& pool) { + if (!pool) { + return Status::Invalid("dense hash table requires a memory pool"); + } + PAIMON_ASSIGN_OR_RAISE(uint64_t bucket_count, BucketCount(capacity, occupancy_pct)); + PAIMON_ASSIGN_OR_RAISE(uint64_t bytes, RequiredBytes(bucket_count)); + auto hash_table = std::unique_ptr<DenseHashTable>(new DenseHashTable()); + hash_table->data_ = AllocateBytesKeepingPoolAlive(bytes, pool); + hash_table->header_ = Header{bucket_count, 0}; + hash_table->capacity_ = capacity; + hash_table->writable_ = true; + hash_table->WriteHeader(); + for (uint64_t index = 0; index < bucket_count; ++index) { + hash_table->WriteBucket(index, Bucket{kEmptyKey, Value{}}); + } + hash_table->WriteBucket(bucket_count, Bucket{0, Value{}}); + hash_table->WriteBucket(bucket_count + 1, Bucket{0, Value{}}); + return hash_table; + } + + /// Mounts a shared buffer without copying. Its bytes must not be mutated while mounted. + /// The returned hash table retains the buffer; its allocator must also outlive the buffer. + static Result<std::unique_ptr<DenseHashTable>> MountForRead(std::shared_ptr<const Bytes> data) { + if (!data || data->size() < kHeaderSize) { + return Status::Invalid("dense hash table buffer is shorter than its header"); + } + Header header; + std::memcpy(&header, data->data(), kHeaderSize); + PAIMON_RETURN_NOT_OK(ValidateHeader(header, data->size())); + auto hash_table = std::unique_ptr<DenseHashTable>(new DenseHashTable()); + hash_table->data_ = std::move(data); + hash_table->header_ = header; + hash_table->capacity_ = header.bucket_count + 2; + for (uint64_t index = header.bucket_count; index < header.bucket_count + 2; ++index) { + if (hash_table->ReadBucket(index).key > 1) { + return Status::Invalid("invalid dense hash table special bucket"); + } + } + return hash_table; + } + + DenseHashTable(const DenseHashTable&) = delete; + DenseHashTable& operator=(const DenseHashTable&) = delete; + + /// New keys fail once capacity is reached. Updating an existing key still succeeds. + Status Insert(Key key, Value value) { + if (!writable_) { + return Status::Invalid("cannot insert into a read-only dense hash table"); + } + if (key == kEmptyKey || key == kOtherSpecialKey) { + const uint64_t index = header_.bucket_count + (key == kOtherSpecialKey); + return SetBucket(index, Bucket{1, value}, ReadBucket(index).key == 0); + } + uint64_t index = key % header_.bucket_count; + for (uint64_t probe = 0; probe < header_.bucket_count; ++probe) { + const Bucket bucket = ReadBucket(index); + if (bucket.key == kEmptyKey || bucket.key == key) { + return SetBucket(index, Bucket{key, value}, bucket.key == kEmptyKey); + } + index = (index + 1) % header_.bucket_count; + } + return Status::Invalid("dense hash table is full"); + } + + std::optional<Value> Find(Key key) const { + if (key == kEmptyKey || key == kOtherSpecialKey) { + const Bucket bucket = ReadBucket(header_.bucket_count + (key == kOtherSpecialKey)); + return bucket.key == 0 ? std::nullopt : std::optional<Value>(bucket.value); + } + uint64_t index = key % header_.bucket_count; + for (uint64_t probe = 0; probe < header_.bucket_count; ++probe) { + const Bucket bucket = ReadBucket(index); + if (bucket.key == kEmptyKey) { + return std::nullopt; + } + if (bucket.key == key) { + return bucket.value; + } + index = (index + 1) % header_.bucket_count; + } + return std::nullopt; + } + + /// Writes the container layout at the stream's current position. The caller owns flush/close. + Status WriteTo(OutputStream* output) const { + if (!output) { + return Status::Invalid("dense hash table output stream is null"); + } + const auto length = static_cast<int64_t>(data_->size()); + PAIMON_ASSIGN_OR_RAISE(int64_t written, output->Write(data_->data(), length)); + if (written != length) { + return Status::IOError("short write for dense hash table"); + } + return Status::OK(); + } + + uint64_t Size() const { + return header_.key_count; + } + + /// Returns the writable key limit, or the physical bucket limit for a read-only mount. + uint64_t Capacity() const { + return capacity_; + } + + uint64_t MemoryUsage() const { + return data_->size(); + } + + private: + friend class DenseHashTableFileReader<Key, Value>; + + struct Header { + uint64_t bucket_count; + uint64_t key_count; + }; + + struct Bucket { + Key key; + Value value; + }; + + static constexpr Key kEmptyKey = static_cast<Key>(0xFFFFFFFFFFFFFEFFULL); + static constexpr Key kOtherSpecialKey = static_cast<Key>(0xFFFFFFFFFFFFFDFFULL); + static constexpr uint64_t kHeaderSize = sizeof(Header); + static constexpr uint64_t kBucketSize = sizeof(Key) + sizeof(Value); + static constexpr uint64_t kMaxBytes = std::numeric_limits<int64_t>::max(); + + DenseHashTable() = default; + + static Result<uint64_t> BucketCount(uint64_t capacity, int32_t occupancy_pct) { + if (occupancy_pct <= 0 || occupancy_pct > 80) { + return Status::Invalid("dense hash table occupancy must be between 1 and 80"); + } + const uint64_t quotient = capacity / occupancy_pct; + if (quotient > kMaxBytes / 100) { + return Status::Invalid("dense hash table capacity overflow"); + } + const uint64_t remainder = + (capacity % occupancy_pct * 100 + occupancy_pct - 1) / occupancy_pct; + return std::max<uint64_t>(1, quotient * 100 + remainder); + } + + static Result<uint64_t> RequiredBytes(uint64_t bucket_count) { + if (bucket_count == 0 || bucket_count > (kMaxBytes - kHeaderSize) / kBucketSize - 2) { + return Status::Invalid("invalid dense hash table bucket count"); + } + return kHeaderSize + (bucket_count + 2) * kBucketSize; + } + + static Status ValidateHeader(const Header& header, uint64_t length) { + PAIMON_ASSIGN_OR_RAISE(uint64_t bytes, RequiredBytes(header.bucket_count)); + if (header.key_count > header.bucket_count + 2 || bytes != length) { + return Status::Invalid( + "dense hash table header does not match its length or key count"); + } + return Status::OK(); + } + + static Bucket DecodeBucket(const char* data) { + Bucket bucket{}; + std::memcpy(&bucket.key, data, sizeof(Key)); + std::memcpy(&bucket.value, data + sizeof(Key), sizeof(Value)); + return bucket; + } + + Bucket ReadBucket(uint64_t index) const { + return DecodeBucket(data_->data() + kHeaderSize + index * kBucketSize); + } + + void WriteBucket(uint64_t index, const Bucket& bucket) { + char* data = data_->data() + kHeaderSize + index * kBucketSize; + std::memcpy(data, &bucket.key, sizeof(Key)); + std::memcpy(data + sizeof(Key), &bucket.value, sizeof(Value)); + } + + void WriteHeader() { + std::memcpy(data_->data(), &header_, kHeaderSize); + } + + Status SetBucket(uint64_t index, const Bucket& bucket, bool is_new) { + if (is_new && Size() >= capacity_) { + return Status::Invalid("dense hash table is full"); + } + WriteBucket(index, bucket); + if (is_new) { + ++header_.key_count; + WriteHeader(); + } + return Status::OK(); + } + + std::shared_ptr<const Bytes> data_; + Header header_{}; + uint64_t capacity_ = 0; + bool writable_ = false; +}; + +} // namespace paimon
diff --git a/src/paimon/common/utils/dense_hash_table_file_reader.h b/src/paimon/common/utils/dense_hash_table_file_reader.h new file mode 100644 index 0000000..2c2ecfe --- /dev/null +++ b/src/paimon/common/utils/dense_hash_table_file_reader.h
@@ -0,0 +1,186 @@ +/* + * 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 <algorithm> +#include <cstdint> +#include <cstring> +#include <limits> +#include <memory> +#include <optional> +#include <string> +#include <utility> + +#include "paimon/common/io/cache/cache_manager.h" +#include "paimon/common/utils/dense_hash_table.h" +#include "paimon/fs/file_system.h" +#include "paimon/macros.h" +#include "paimon/memory/memory_segment.h" +#include "paimon/result.h" + +namespace paimon { + +/// Reads the native DenseHashTable layout using positional reads and a shared CacheManager. +/// Opening reads only the header, not the bucket array. Find is safe for concurrent callers if +/// the input stream supports concurrent positional reads. Shared streams are not closed here. +/// File paths identify immutable contents and must be unique within the shared cache. Destroying +/// a reader does not invalidate cached pages, which retain their allocation's memory pool. +template <typename Key, typename Value> +class DenseHashTableFileReader { + public: + static Result<std::unique_ptr<DenseHashTableFileReader>> Create( + uint32_t page_size, const std::shared_ptr<CacheManager>& cache_manager, + const std::shared_ptr<FileSystem>& fs, const std::string& path, + const std::shared_ptr<MemoryPool>& pool) { + if (!fs) { + return Status::Invalid("dense hash table file system is null"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input, fs->Open(path)); + return Create(page_size, cache_manager, path, input, pool); + } + + static Result<std::unique_ptr<DenseHashTableFileReader>> Create( + uint32_t page_size, const std::shared_ptr<CacheManager>& cache_manager, + const std::string& file_path, const std::shared_ptr<InputStream>& input, + const std::shared_ptr<MemoryPool>& pool) { + if (!input || !pool || !cache_manager || file_path.empty() || page_size == 0 || + page_size > static_cast<uint32_t>(std::numeric_limits<int32_t>::max())) { + return Status::Invalid("invalid dense hash table reader input, pool or cache options"); + } + PAIMON_ASSIGN_OR_RAISE(int64_t length, input->Length()); + if (length < static_cast<int64_t>(HashTable::kHeaderSize)) { + return Status::Invalid("dense hash table file is shorter than its header"); + } + typename HashTable::Header header; + PAIMON_RETURN_NOT_OK( + ReadAt(input.get(), reinterpret_cast<char*>(&header), HashTable::kHeaderSize, 0)); + PAIMON_RETURN_NOT_OK(HashTable::ValidateHeader(header, length)); + return std::unique_ptr<DenseHashTableFileReader>(new DenseHashTableFileReader( + page_size, cache_manager, file_path, input, header, length, pool)); + } + + DenseHashTableFileReader(const DenseHashTableFileReader&) = delete; + DenseHashTableFileReader& operator=(const DenseHashTableFileReader&) = delete; + + Result<std::optional<Value>> Find(Key key) const { + std::optional<Page> page; + if (key == HashTable::kEmptyKey || key == HashTable::kOtherSpecialKey) { + PAIMON_ASSIGN_OR_RAISE( + typename HashTable::Bucket bucket, + ReadBucket(header_.bucket_count + (key == HashTable::kOtherSpecialKey), &page)); + if (bucket.key > 1) { + return Status::Invalid("invalid dense hash table special bucket"); + } + return bucket.key == 0 ? std::nullopt : std::optional<Value>(bucket.value); + } + uint64_t index = key % header_.bucket_count; + for (uint64_t probe = 0; probe < header_.bucket_count; ++probe) { + PAIMON_ASSIGN_OR_RAISE(typename HashTable::Bucket bucket, ReadBucket(index, &page)); + if (bucket.key == HashTable::kEmptyKey) { + return std::optional<Value>(); + } + if (bucket.key == HashTable::kOtherSpecialKey) { + return Status::Invalid("invalid dense hash table ordinary bucket"); + } + if (bucket.key == key) { + return std::optional<Value>(bucket.value); + } + index = (index + 1) % header_.bucket_count; + } + return std::optional<Value>(); + } + + uint64_t Size() const { + return header_.key_count; + } + + private: + using HashTable = DenseHashTable<Key, Value>; + + struct Page { + uint64_t index; + MemorySegment segment; + }; + + DenseHashTableFileReader(uint32_t page_size, std::shared_ptr<CacheManager> cache_manager, + std::string file_path, std::shared_ptr<InputStream> input, + typename HashTable::Header header, uint64_t length, + std::shared_ptr<MemoryPool> pool) + : input_(std::move(input)), + pool_(std::move(pool)), + cache_manager_(std::move(cache_manager)), + file_path_(std::move(file_path)), + header_(header), + length_(length), + page_size_(page_size) {} + + static Status ReadAt(InputStream* input, char* data, int64_t size, int64_t offset) { + PAIMON_ASSIGN_OR_RAISE(int64_t read, input->Read(data, size, offset)); + if (read != size) { + return Status::IOError("short read for dense hash table"); + } + return Status::OK(); + } + + Result<MemorySegment> ReadPage(uint64_t index) const { + const uint64_t offset = index * page_size_; + const int32_t size = static_cast<int32_t>(std::min<uint64_t>(page_size_, length_ - offset)); + std::shared_ptr<CacheKey> key = + CacheKey::ForPosition(file_path_, offset, size, /*is_index=*/true); + return cache_manager_->GetPage( + key, + [this, offset, size](const std::shared_ptr<CacheKey>&) -> Result<MemorySegment> { + std::shared_ptr<Bytes> bytes = AllocateBytesKeepingPoolAlive(size, pool_); + PAIMON_RETURN_NOT_OK(ReadAt(input_.get(), bytes->data(), size, offset)); + return MemorySegment::Wrap(bytes); + }, + /*eviction_callback=*/nullptr); + } + + Result<typename HashTable::Bucket> ReadBucket(uint64_t index, std::optional<Page>* page) const { + uint64_t offset = HashTable::kHeaderSize + index * HashTable::kBucketSize; + char data[HashTable::kBucketSize]; + uint64_t copied = 0; + while (copied < HashTable::kBucketSize) { + const uint64_t page_index = offset / page_size_; + if (!*page || (*page)->index != page_index) { + PAIMON_ASSIGN_OR_RAISE(MemorySegment segment, ReadPage(page_index)); + *page = Page{page_index, std::move(segment)}; + } + const uint64_t position = offset % page_size_; + const uint64_t size = std::min<uint64_t>((*page)->segment.Size() - position, + HashTable::kBucketSize - copied); + std::memcpy(data + copied, (*page)->segment.Data() + position, size); + copied += size; + offset += size; + } + return HashTable::DecodeBucket(data); + } + + std::shared_ptr<InputStream> input_; + std::shared_ptr<MemoryPool> pool_; + std::shared_ptr<CacheManager> cache_manager_; + std::string file_path_; + typename HashTable::Header header_; + uint64_t length_; + uint32_t page_size_; +}; + +} // namespace paimon
diff --git a/src/paimon/common/utils/dense_hash_table_file_reader_test.cpp b/src/paimon/common/utils/dense_hash_table_file_reader_test.cpp new file mode 100644 index 0000000..a301cd0 --- /dev/null +++ b/src/paimon/common/utils/dense_hash_table_file_reader_test.cpp
@@ -0,0 +1,448 @@ +/* + * 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/common/utils/dense_hash_table_file_reader.h" + +#include <atomic> +#include <cstdint> +#include <cstring> +#include <limits> +#include <memory> +#include <optional> +#include <string> +#include <thread> +#include <utility> +#include <vector> + +#include "gtest/gtest.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/cache/cache_manager.h" +#include "paimon/common/io/cache/lru_cache.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/memory/bytes_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/dense_hash_table.h" +#include "paimon/common/utils/hash_map.h" +#include "paimon/fs/file_system.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +using HashTable = DenseHashTable<uint64_t, int64_t>; +using Reader = DenseHashTableFileReader<uint64_t, int64_t>; + +int64_t CachedBytes(const std::shared_ptr<CacheManager>& cache_manager) { + return checked_pointer_cast<LruCache>(cache_manager->IndexCache())->GetCurrentWeight(); +} + +class CountingInputStream : public ByteArrayInputStream { + public: + explicit CountingInputStream(std::shared_ptr<Bytes> bytes) + : ByteArrayInputStream(bytes->data(), bytes->size()), bytes_(std::move(bytes)) {} + + using ByteArrayInputStream::Read; + + Result<int64_t> Read(char* buffer, int64_t size, int64_t offset) override { + ++read_calls; + if (fail_reads) { + return Status::IOError("injected dense hash table read failure"); + } + PAIMON_ASSIGN_OR_RAISE(int64_t read, ByteArrayInputStream::Read( + buffer, short_reads ? size - 1 : size, offset)); + bytes_read += read; + return read; + } + + std::atomic<int64_t> read_calls{0}; + std::atomic<int64_t> bytes_read{0}; + bool short_reads = false; + bool fail_reads = false; + + private: + std::shared_ptr<Bytes> bytes_; +}; + +Result<std::shared_ptr<Bytes>> FlushHashTable(const HashTable& hash_table, + const std::shared_ptr<MemoryPool>& pool) { + ByteArrayOutputStream output(std::make_unique<MemorySegmentOutputStream>(1024, pool)); + PAIMON_RETURN_NOT_OK(hash_table.WriteTo(&output)); + PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> bytes, output.Finish(pool.get())); + std::shared_ptr<Bytes> result = AllocateBytesKeepingPoolAlive(bytes->size(), pool); + std::memcpy(result->data(), bytes->data(), bytes->size()); + return result; +} + +Result<std::shared_ptr<Bytes>> MakeHashTableBytes(const std::shared_ptr<MemoryPool>& pool) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<HashTable> hash_table, + HashTable::Create(1000, /*occupancy_pct=*/65, pool)); + for (uint64_t key = 0; key < 998; ++key) { + PAIMON_RETURN_NOT_OK(hash_table->Insert(key, key + 10)); + } + PAIMON_RETURN_NOT_OK(hash_table->Insert(0xFFFFFFFFFFFFFEFFULL, -1)); + PAIMON_RETURN_NOT_OK( + hash_table->Insert(0xFFFFFFFFFFFFFDFFULL, std::numeric_limits<int64_t>::min())); + return FlushHashTable(*hash_table, pool); +} + +} // namespace + +TEST(DenseHashTableFileReaderTest, TestOpenIsLazyAndRepeatedLookupsUseCache) { + auto cache_manager = std::make_shared<CacheManager>(512, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Reader> reader, + Reader::Create(/*page_size=*/256, cache_manager, "dense-hash-table", input, + GetDefaultPool())); + ASSERT_EQ(reader->Size(), 1000); + ASSERT_EQ(input->read_calls.load(), 1); + ASSERT_EQ(input->bytes_read.load(), 16); + ASSERT_EQ(CachedBytes(cache_manager), 0); + + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(input->read_calls.load(), 2); + ASSERT_EQ(input->bytes_read.load(), 16 + 256); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(input->read_calls.load(), 2); + ASSERT_EQ(CachedBytes(cache_manager), 256); + ASSERT_OK_AND_ASSIGN(int64_t position, input->GetPos()); + ASSERT_EQ(position, 0); +} + +TEST(DenseHashTableFileReaderTest, TestPageBoundariesSpecialKeysAndMissingKeys) { + auto cache_manager = std::make_shared<CacheManager>(4 * 17, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + // Seventeen-byte pages force buckets to span page boundaries. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(17, cache_manager, "dense-hash-table", input, GetDefaultPool())); + for (uint64_t key : {uint64_t{0}, uint64_t{1}, uint64_t{30}, uint64_t{997}}) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(key)); + ASSERT_EQ(value, static_cast<int64_t>(key + 10)); + ASSERT_LE(CachedBytes(cache_manager), 4 * 17); + } + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(0xFFFFFFFFFFFFFEFFULL)); + ASSERT_EQ(value, -1); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0xFFFFFFFFFFFFFDFFULL)); + ASSERT_EQ(value, std::numeric_limits<int64_t>::min()); + ASSERT_OK_AND_ASSIGN(value, reader->Find(1000)); + ASSERT_FALSE(value); + ASSERT_OK_AND_ASSIGN(value, reader->Find(std::numeric_limits<uint64_t>::max())); + ASSERT_FALSE(value); +} + +TEST(DenseHashTableFileReaderTest, TestProbingAcrossPagesAndWraparound) { + auto cache_manager = std::make_shared<CacheManager>(32, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(3, /*occupancy_pct=*/50, GetDefaultPool())); + ASSERT_OK(hash_table->Insert(5, 1)); + ASSERT_OK(hash_table->Insert(11, 2)); + ASSERT_OK(hash_table->Insert(17, 3)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, + FlushHashTable(*hash_table, GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(16, cache_manager, "dense-hash-table", input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(17)); + ASSERT_EQ(value, 3); + ASSERT_OK_AND_ASSIGN(value, reader->Find(23)); + ASSERT_FALSE(value); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0xFFFFFFFFFFFFFEFFULL)); + ASSERT_FALSE(value); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0xFFFFFFFFFFFFFDFFULL)); + ASSERT_FALSE(value); +} + +TEST(DenseHashTableFileReaderTest, TestEmptyHashTable) { + auto cache_manager = + std::make_shared<CacheManager>(8 * 1024 * 1024, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(0, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, + FlushHashTable(*hash_table, GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(64 * 1024, cache_manager, "dense-hash-table", input, GetDefaultPool())); + ASSERT_EQ(reader->Size(), 0); + for (uint64_t key : + {uint64_t{0}, uint64_t{0xFFFFFFFFFFFFFEFFULL}, uint64_t{0xFFFFFFFFFFFFFDFFULL}}) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(key)); + ASSERT_FALSE(value); + } +} + +TEST(DenseHashTableFileReaderTest, TestMixedIntegerWidths) { + auto cache_manager = std::make_shared<CacheManager>(68, /*high_priority_pool_ratio=*/0.0); + using MixedHashTable = DenseHashTable<uint32_t, int64_t>; + using MixedReader = DenseHashTableFileReader<uint32_t, int64_t>; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<MixedHashTable> hash_table, + MixedHashTable::Create(3, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_OK(hash_table->Insert(1, -10)); + ASSERT_OK(hash_table->Insert(0xFFFFFEFFU, 20)); + ASSERT_OK(hash_table->Insert(0xFFFFFDFFU, 30)); + ByteArrayOutputStream output( + std::make_unique<MemorySegmentOutputStream>(1024, GetDefaultPool())); + ASSERT_OK(hash_table->WriteTo(&output)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, output.Finish(GetDefaultPool().get())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<MixedReader> reader, + MixedReader::Create(17, cache_manager, "dense-hash-table", input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(1)); + ASSERT_EQ(value, -10); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0xFFFFFEFFU)); + ASSERT_EQ(value, 20); + ASSERT_OK_AND_ASSIGN(value, reader->Find(0xFFFFFDFFU)); + ASSERT_EQ(value, 30); +} + +TEST(DenseHashTableFileReaderTest, TestCacheEvictionAndDisabledCache) { + auto cache_manager = std::make_shared<CacheManager>(256, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool())); + for (uint64_t key : {uint64_t{0}, uint64_t{500}, uint64_t{0}}) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(key)); + ASSERT_EQ(value, static_cast<int64_t>(key + 10)); + ASSERT_EQ(CachedBytes(cache_manager), 256); + } + ASSERT_EQ(input->read_calls.load(), 4); + reader.reset(); + cache_manager = std::make_shared<CacheManager>(0, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN( + reader, Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool())); + const int64_t before = input->read_calls.load(); + for (int32_t iteration = 0; iteration < 2; ++iteration) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(CachedBytes(cache_manager), 0); + } + ASSERT_EQ(input->read_calls.load(), before + 2); +} + +TEST(DenseHashTableFileReaderTest, TestSharedCacheSurvivesReaderDestruction) { + auto cache_manager = std::make_shared<CacheManager>(512, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto first_input = std::make_shared<CountingInputStream>(bytes); + auto second_input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> first, + Reader::Create(256, cache_manager, "same-file", first_input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> second, + Reader::Create(256, cache_manager, "same-file", second_input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, first->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(first_input->read_calls.load(), 2); + first.reset(); + ASSERT_EQ(CachedBytes(cache_manager), 256); + + ASSERT_OK_AND_ASSIGN(value, second->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(second_input->read_calls.load(), 1); + second.reset(); + ASSERT_OK_AND_ASSIGN( + second, Reader::Create(256, cache_manager, "same-file", second_input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(value, second->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(second_input->read_calls.load(), 2); +} + +TEST(DenseHashTableFileReaderTest, TestSharedBudgetAndDistinctFiles) { + auto cache_manager = std::make_shared<CacheManager>(256, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> first_bytes, MakeHashTableBytes(GetDefaultPool())); + std::shared_ptr<Bytes> second_bytes = + AllocateBytesKeepingPoolAlive(first_bytes->size(), GetDefaultPool()); + std::memcpy(second_bytes->data(), first_bytes->data(), first_bytes->size()); + // Change key zero's value in the second file; both lookups read the same byte range. + const int64_t second_value = 1000; + std::memcpy(second_bytes->data() + 16 + sizeof(uint64_t), &second_value, sizeof(second_value)); + auto first_input = std::make_shared<CountingInputStream>(first_bytes); + auto second_input = std::make_shared<CountingInputStream>(second_bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> first, + Reader::Create(256, cache_manager, "first-file", first_input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> second, + Reader::Create(256, cache_manager, "second-file", second_input, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, first->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_OK_AND_ASSIGN(value, second->Find(0)); + ASSERT_EQ(value, second_value); + ASSERT_EQ(CachedBytes(cache_manager), 256); + ASSERT_EQ(cache_manager->IndexCache()->Size(), 1); + ASSERT_OK_AND_ASSIGN(value, first->Find(0)); + ASSERT_EQ(value, 10); + ASSERT_EQ(first_input->read_calls.load(), 3); + ASSERT_EQ(second_input->read_calls.load(), 2); + ASSERT_EQ(CachedBytes(cache_manager), 256); +} + +TEST(DenseHashTableFileReaderTest, TestHashMapToDenseFileRoundTrip) { + auto cache_manager = + std::make_shared<CacheManager>(8 * 1024 * 1024, /*high_priority_pool_ratio=*/0.0); + std::unique_ptr<UniqueTestDirectory> directory = UniqueTestDirectory::Create(); + ASSERT_NE(directory, nullptr); + using Map = HashMap<uint64_t, int64_t>; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Map> map, Map::Create(1, GetDefaultPool())); + for (uint64_t key = 0; key < 300; ++key) { + ASSERT_OK(map->FindAndInsert(key, key)); + } + ASSERT_OK(map->FindAndInsert(0, -1)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(map->Size(), /*occupancy_pct=*/65, GetDefaultPool())); + Status inserted; + map->ForEach([&](uint64_t key, int64_t value) { + if (inserted.ok()) { + inserted = hash_table->Insert(key, value); + } + }); + ASSERT_OK(inserted); + const std::string path = directory->Str() + "/dense-hash-table"; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> output, + directory->GetFileSystem()->Create(path, /*overwrite=*/false)); + ASSERT_OK(hash_table->WriteTo(output.get())); + ASSERT_OK(output->Flush()); + ASSERT_OK(output->Close()); + hash_table.reset(); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Reader> reader, + Reader::Create(64 * 1024, cache_manager, directory->GetFileSystem(), path, + GetDefaultPool())); + for (uint64_t key = 0; key < 300; ++key) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(key)); + ASSERT_EQ(value, map->Find(key)); + } +} + +TEST(DenseHashTableFileReaderTest, TestConcurrentLookupsAndBufferLifetime) { + auto cache_manager = std::make_shared<CacheManager>(512, /*high_priority_pool_ratio=*/0.0); + std::shared_ptr<MemoryPool> pool = GetMemoryPool(); + std::weak_ptr<MemoryPool> weak_pool = pool; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(pool)); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Reader> reader, + Reader::Create(256, cache_manager, "dense-hash-table", input, pool)); + bytes.reset(); + input.reset(); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + std::vector<std::thread> threads; + for (int32_t thread = 0; thread < 4; ++thread) { + threads.emplace_back([&reader]() { + for (uint64_t key = 0; key < 998; ++key) { + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(key)); + ASSERT_EQ(value, static_cast<int64_t>(key + 10)); + } + }); + } + for (std::thread& thread : threads) { + thread.join(); + } + ASSERT_LE(CachedBytes(cache_manager), 512); + reader.reset(); + ASSERT_FALSE(weak_pool.expired()); + cache_manager.reset(); + ASSERT_TRUE(weak_pool.expired()); +} + +TEST(DenseHashTableFileReaderTest, TestInvalidOptionsAndHeader) { + auto cache_manager = std::make_shared<CacheManager>(512, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_NOK(Reader::Create(256, cache_manager, "dense-hash-table", + std::shared_ptr<InputStream>(), GetDefaultPool())); + ASSERT_NOK(Reader::Create(256, cache_manager, "dense-hash-table", input, nullptr)); + ASSERT_NOK(Reader::Create(0, cache_manager, "dense-hash-table", input, GetDefaultPool())); + ASSERT_NOK(Reader::Create(std::numeric_limits<uint32_t>::max(), cache_manager, + "dense-hash-table", input, GetDefaultPool())); + ASSERT_NOK(Reader::Create(16, nullptr, "dense-hash-table", input, GetDefaultPool())); + ASSERT_NOK(Reader::Create(256, cache_manager, "", input, GetDefaultPool())); + ASSERT_NOK(Reader::Create(256, cache_manager, std::shared_ptr<FileSystem>(), "unused", + GetDefaultPool())); + + auto short_input = + std::make_shared<CountingInputStream>(AllocateBytesKeepingPoolAlive(15, GetDefaultPool())); + ASSERT_NOK_WITH_MSG( + Reader::Create(256, cache_manager, "dense-hash-table", short_input, GetDefaultPool()), + "shorter than its header"); + const uint64_t zero = 0; + std::memcpy(bytes->data(), &zero, sizeof(zero)); + ASSERT_NOK_WITH_MSG( + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool()), + "bucket count"); + const uint64_t overflow = std::numeric_limits<uint64_t>::max(); + std::memcpy(bytes->data(), &overflow, sizeof(overflow)); + ASSERT_NOK_WITH_MSG( + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool()), + "bucket count"); +} + +TEST(DenseHashTableFileReaderTest, TestReadErrorsArePropagatedAndNotCached) { + auto cache_manager = std::make_shared<CacheManager>(512, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + auto input = std::make_shared<CountingInputStream>(bytes); + input->short_reads = true; + ASSERT_NOK_WITH_MSG( + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool()), + "short read"); + input->short_reads = false; + input->fail_reads = true; + ASSERT_NOK_WITH_MSG( + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool()), + "injected dense hash table read failure"); + input->fail_reads = false; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(256, cache_manager, "dense-hash-table", input, GetDefaultPool())); + input->short_reads = true; + ASSERT_NOK_WITH_MSG(reader->Find(0), "short read"); + ASSERT_EQ(CachedBytes(cache_manager), 0); + input->short_reads = false; + input->fail_reads = true; + ASSERT_NOK_WITH_MSG(reader->Find(0), "injected dense hash table read failure"); + ASSERT_EQ(CachedBytes(cache_manager), 0); + input->fail_reads = false; + ASSERT_OK_AND_ASSIGN(std::optional<int64_t> value, reader->Find(0)); + ASSERT_EQ(value, 10); +} + +TEST(DenseHashTableFileReaderTest, TestRejectInvalidSpecialBucket) { + auto cache_manager = + std::make_shared<CacheManager>(8 * 1024 * 1024, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, MakeHashTableBytes(GetDefaultPool())); + const uint64_t invalid_state = 2; + std::memcpy(bytes->data() + bytes->size() - 32, &invalid_state, sizeof(invalid_state)); + auto input = std::make_shared<CountingInputStream>(bytes); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<Reader> reader, + Reader::Create(64 * 1024, cache_manager, "dense-hash-table", input, GetDefaultPool())); + ASSERT_NOK_WITH_MSG(reader->Find(0xFFFFFFFFFFFFFEFFULL), "special bucket"); +} + +} // namespace paimon::test
diff --git a/src/paimon/common/utils/dense_hash_table_test.cpp b/src/paimon/common/utils/dense_hash_table_test.cpp new file mode 100644 index 0000000..75bcfb5 --- /dev/null +++ b/src/paimon/common/utils/dense_hash_table_test.cpp
@@ -0,0 +1,195 @@ +/* + * 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/common/utils/dense_hash_table.h" + +#include <cstdint> +#include <cstring> +#include <limits> +#include <memory> +#include <vector> + +#include "gtest/gtest.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/memory/bytes_utils.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/mock/mock_file_system.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +using HashTable = DenseHashTable<uint64_t, int64_t>; + +template <typename Key, typename Value> +Result<std::shared_ptr<Bytes>> FlushHashTable(const DenseHashTable<Key, Value>& hash_table, + const std::shared_ptr<MemoryPool>& pool) { + ByteArrayOutputStream output(std::make_unique<MemorySegmentOutputStream>(1024, pool)); + PAIMON_RETURN_NOT_OK(hash_table.WriteTo(&output)); + PAIMON_ASSIGN_OR_RAISE(PAIMON_UNIQUE_PTR<Bytes> bytes, output.Finish(pool.get())); + std::shared_ptr<Bytes> result = AllocateBytesKeepingPoolAlive(bytes->size(), pool); + std::memcpy(result->data(), bytes->data(), bytes->size()); + return result; +} + +class FailingOutputStream : public MockOutputStream { + public: + Result<int64_t> Write(const char*, int64_t) override { + return Status::IOError("injected dense hash table write failure"); + } +}; + +} // namespace + +TEST(DenseHashTableTest, TestProbingWraparoundAndCapacity) { + // Three entries at 50% occupancy give six buckets. These keys all start at bucket five. + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(3, /*occupancy_pct=*/50, GetDefaultPool())); + ASSERT_EQ(hash_table->Capacity(), 3); + ASSERT_OK(hash_table->Insert(5, 10)); + ASSERT_OK(hash_table->Insert(11, 20)); + ASSERT_OK(hash_table->Insert(17, 30)); + ASSERT_EQ(hash_table->Find(5), 10); + ASSERT_EQ(hash_table->Find(11), 20); + ASSERT_EQ(hash_table->Find(17), 30); + ASSERT_FALSE(hash_table->Find(23)); + ASSERT_NOK_WITH_MSG(hash_table->Insert(23, 40), "full"); + ASSERT_OK(hash_table->Insert(5, -1)); + ASSERT_EQ(hash_table->Find(5), -1); + ASSERT_EQ(hash_table->Size(), 3); +} + +TEST(DenseHashTableTest, TestAllKeyAndValueBitPatterns) { + const std::vector<uint64_t> keys = {0, + 1, + 0xFFFFFFFFFFFFFEFFULL, + 0xFFFFFFFFFFFFFDFFULL, + std::numeric_limits<uint64_t>::max(), + uint64_t{1} << 63}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(keys.size(), /*occupancy_pct=*/65, GetDefaultPool())); + for (size_t index = 0; index < keys.size(); ++index) { + ASSERT_FALSE(hash_table->Find(keys[index])); + ASSERT_OK(hash_table->Insert(keys[index], static_cast<int64_t>(index) - 3)); + } + for (size_t index = 0; index < keys.size(); ++index) { + ASSERT_EQ(hash_table->Find(keys[index]), static_cast<int64_t>(index) - 3); + } + ASSERT_OK(hash_table->Insert(keys[2], std::numeric_limits<int64_t>::min())); + ASSERT_OK(hash_table->Insert(keys[3], std::numeric_limits<int64_t>::max())); + ASSERT_EQ(hash_table->Find(keys[2]), std::numeric_limits<int64_t>::min()); + ASSERT_EQ(hash_table->Find(keys[3]), std::numeric_limits<int64_t>::max()); + ASSERT_EQ(hash_table->Size(), keys.size()); + ASSERT_NOK(hash_table->Insert(2, 4)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> full, + HashTable::Create(1, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_OK(full->Insert(1, 1)); + ASSERT_NOK(full->Insert(keys[2], 2)); + ASSERT_NOK(full->Insert(keys[3], 3)); + ASSERT_FALSE(full->Find(keys[2])); + ASSERT_FALSE(full->Find(keys[3])); +} + +TEST(DenseHashTableTest, TestEmptyAndInvalidArguments) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(0, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_EQ(hash_table->Size(), 0); + ASSERT_FALSE(hash_table->Find(0)); + ASSERT_FALSE(hash_table->Find(0xFFFFFFFFFFFFFEFFULL)); + ASSERT_FALSE(hash_table->Find(0xFFFFFFFFFFFFFDFFULL)); + ASSERT_NOK_WITH_MSG(hash_table->Insert(0, 0), "full"); + ASSERT_NOK(HashTable::Create(1, 65, nullptr)); + ASSERT_NOK(HashTable::Create(1, 0, GetDefaultPool())); + ASSERT_NOK(HashTable::Create(1, -1, GetDefaultPool())); + ASSERT_NOK(HashTable::Create(1, 81, GetDefaultPool())); + ASSERT_NOK(HashTable::Create(std::numeric_limits<uint64_t>::max(), 65, GetDefaultPool())); + ASSERT_NOK(HashTable::Create(std::numeric_limits<int64_t>::max(), 80, GetDefaultPool())); +} + +TEST(DenseHashTableTest, TestFlushAndReadOnlyMountKeepPoolAlive) { + std::shared_ptr<MemoryPool> pool = GetMemoryPool(); + std::weak_ptr<MemoryPool> weak_pool = pool; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(5000, /*occupancy_pct=*/65, pool)); + for (uint64_t key = 0; key < 5000; ++key) { + ASSERT_OK(hash_table->Insert(key, key + 10)); + } + ASSERT_EQ(pool->CurrentUsage(), hash_table->MemoryUsage()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, FlushHashTable(*hash_table, pool)); + ASSERT_EQ(bytes->size(), hash_table->MemoryUsage()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> mounted, HashTable::MountForRead(bytes)); + hash_table.reset(); + bytes.reset(); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + ASSERT_EQ(mounted->Size(), 5000); + for (uint64_t key = 0; key < 5000; ++key) { + ASSERT_EQ(mounted->Find(key), static_cast<int64_t>(key + 10)); + } + ASSERT_FALSE(mounted->Find(5000)); + ASSERT_NOK_WITH_MSG(mounted->Insert(0, 0), "read-only"); + mounted.reset(); + ASSERT_TRUE(weak_pool.expired()); +} + +TEST(DenseHashTableTest, TestRejectMalformedBuffers) { + ASSERT_NOK(HashTable::MountForRead(nullptr)); + ASSERT_NOK(HashTable::MountForRead(AllocateBytesKeepingPoolAlive(15, GetDefaultPool()))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(1, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, + FlushHashTable(*hash_table, GetDefaultPool())); + + const uint64_t zero = 0; + const uint64_t overflow = std::numeric_limits<uint64_t>::max(); + std::memcpy(bytes->data(), &zero, sizeof(zero)); + ASSERT_NOK_WITH_MSG(HashTable::MountForRead(bytes), "bucket count"); + std::memcpy(bytes->data(), &overflow, sizeof(overflow)); + ASSERT_NOK_WITH_MSG(HashTable::MountForRead(bytes), "bucket count"); + + ASSERT_OK_AND_ASSIGN(bytes, FlushHashTable(*hash_table, GetDefaultPool())); + std::memcpy(bytes->data() + sizeof(uint64_t), &overflow, sizeof(overflow)); + ASSERT_NOK_WITH_MSG(HashTable::MountForRead(bytes), "key count"); + + ASSERT_OK_AND_ASSIGN(bytes, FlushHashTable(*hash_table, GetDefaultPool())); + const uint64_t invalid_state = 2; + std::memcpy(bytes->data() + bytes->size() - 32, &invalid_state, sizeof(invalid_state)); + ASSERT_NOK_WITH_MSG(HashTable::MountForRead(bytes), "special bucket"); + + ASSERT_OK_AND_ASSIGN(bytes, FlushHashTable(*hash_table, GetDefaultPool())); + std::shared_ptr<Bytes> truncated = + AllocateBytesKeepingPoolAlive(bytes->size() - 1, GetDefaultPool()); + std::memcpy(truncated->data(), bytes->data(), truncated->size()); + ASSERT_NOK_WITH_MSG(HashTable::MountForRead(truncated), "length"); +} + +TEST(DenseHashTableTest, TestWriteErrors) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(1, /*occupancy_pct=*/65, GetDefaultPool())); + ASSERT_NOK(hash_table->WriteTo(nullptr)); + MockOutputStream short_output; + ASSERT_NOK_WITH_MSG(hash_table->WriteTo(&short_output), "short write"); + FailingOutputStream failed_output; + ASSERT_NOK_WITH_MSG(hash_table->WriteTo(&failed_output), + "injected dense hash table write failure"); +} + +} // namespace paimon::test
diff --git a/src/paimon/common/utils/hash_map.h b/src/paimon/common/utils/hash_map.h new file mode 100644 index 0000000..b6da896 --- /dev/null +++ b/src/paimon/common/utils/hash_map.h
@@ -0,0 +1,206 @@ +/* + * Copyright 2014-present Alibaba Inc. + * + * Licensed 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. + */ + +// Adapted from Alibaba Havenask +// https://github.com/alibaba/havenask/blob/main/aios/storage/indexlib/util/HashMap.h + +#pragma once + +#include <cstddef> +#include <cstdint> +#include <cstring> +#include <limits> +#include <memory> +#include <optional> +#include <type_traits> +#include <utility> +#include <vector> + +#include "paimon/common/memory/bytes_utils.h" +#include "paimon/macros.h" +#include "paimon/result.h" + +namespace paimon { + +/// A pool-backed hash map which grows by adding blocks without rehashing existing entries. +/// Keys are already-hashed uint32_t or uint64_t values. Values must be trivially copyable and +/// default constructible. Mutation requires external synchronization; immutable maps may be shared. +template <typename Key, typename Value> +class HashMap { + static_assert(std::is_same_v<Key, uint32_t> || std::is_same_v<Key, uint64_t>); + static_assert(std::is_trivially_copyable_v<Value> && std::is_default_constructible_v<Value>); + + public: + static Result<std::unique_ptr<HashMap>> Create(uint64_t initial_capacity, + const std::shared_ptr<MemoryPool>& pool) { + if (!pool || initial_capacity == 0 || initial_capacity > kMaxBytes / 2) { + return Status::Invalid("hash map requires a pool and a valid positive capacity"); + } + auto map = std::unique_ptr<HashMap>(new HashMap(pool)); + PAIMON_RETURN_NOT_OK(map->AddBlock(initial_capacity * 2)); + return map; + } + + HashMap(const HashMap&) = delete; + HashMap& operator=(const HashMap&) = delete; + + std::optional<Value> Find(Key key) const { + for (auto iter = blocks_.rbegin(); iter != blocks_.rend(); ++iter) { + const Block& block = *iter; + uint64_t position = key % block.bucket_count; + for (uint64_t probe = 0; probe < block.bucket_count; ++probe) { + if (!block.Occupied(position)) { + break; + } + if (block.ReadKey(position) == key) { + return block.ReadValue(position); + } + position = (position + 1) % block.bucket_count; + } + } + return std::nullopt; + } + + /// Inserts or updates a key across all blocks. Each key is stored exactly once. + Status FindAndInsert(Key key, const Value& value) { + for (auto iter = blocks_.rbegin(); iter != blocks_.rend(); ++iter) { + Block& block = *iter; + uint64_t position = key % block.bucket_count; + for (uint64_t probe = 0; probe < block.bucket_count; ++probe) { + if (!block.Occupied(position)) { + break; + } + if (block.ReadKey(position) == key) { + block.WriteValue(position, value); + return Status::OK(); + } + position = (position + 1) % block.bucket_count; + } + } + if (blocks_.empty()) { + PAIMON_RETURN_NOT_OK(AddBlock(initial_bucket_count_)); + } else if (blocks_.back().size >= blocks_.back().bucket_count / 2) { + if (blocks_.back().bucket_count > kMaxBytes / 2) { + return Status::Invalid("hash map capacity overflow"); + } + PAIMON_RETURN_NOT_OK(AddBlock(blocks_.back().bucket_count * 2)); + } + Block& block = blocks_.back(); + uint64_t position = key % block.bucket_count; + while (block.Occupied(position)) { + position = (position + 1) % block.bucket_count; + } + block.WriteValue(position, value); + std::memcpy(block.Entry(position), &key, sizeof(Key)); + block.bitmap[position / 8] |= static_cast<uint8_t>(1U << (position % 8)); + ++block.size; + ++size_; + return Status::OK(); + } + + /// Visits each key/value pair once. The callback must not mutate this map. + template <typename Visitor> + void ForEach(Visitor&& visitor) const { + for (const Block& block : blocks_) { + for (uint64_t position = 0; position < block.bucket_count; ++position) { + if (block.Occupied(position)) { + visitor(block.ReadKey(position), block.ReadValue(position)); + } + } + } + } + + /// Releases all bucket allocations. The next insertion reuses the initial capacity. + void Clear() { + blocks_.clear(); + size_ = 0; + memory_usage_ = 0; + } + + uint64_t Size() const { + return size_; + } + + /// Returns bucket and occupancy-bitmap bytes allocated from the pool. + uint64_t MemoryUsage() const { + return memory_usage_; + } + + private: + static constexpr uint64_t kEntrySize = sizeof(Key) + sizeof(Value); + static constexpr uint64_t kMaxBytes = std::numeric_limits<int64_t>::max(); + + struct Block { + char* Entry(uint64_t position) const { + return data->data() + position * kEntrySize; + } + + bool Occupied(uint64_t position) const { + return (bitmap[position / 8] & (1U << (position % 8))) != 0; + } + + Key ReadKey(uint64_t position) const { + Key key{}; + std::memcpy(&key, Entry(position), sizeof(Key)); + return key; + } + + Value ReadValue(uint64_t position) const { + Value value{}; + std::memcpy(&value, Entry(position) + sizeof(Key), sizeof(Value)); + return value; + } + + void WriteValue(uint64_t position, const Value& value) { + std::memcpy(Entry(position) + sizeof(Key), &value, sizeof(Value)); + } + + std::shared_ptr<Bytes> data; + uint8_t* bitmap = nullptr; + uint64_t bucket_count = 0; + uint64_t size = 0; + }; + + explicit HashMap(std::shared_ptr<MemoryPool> pool) : pool_(std::move(pool)) {} + + Status AddBlock(uint64_t bucket_count) { + const uint64_t bitmap_size = bucket_count / 8 + (bucket_count % 8 != 0); + if (bucket_count > (kMaxBytes - bitmap_size) / kEntrySize) { + return Status::Invalid("hash map allocation size overflow"); + } + const uint64_t bytes = bucket_count * kEntrySize + bitmap_size; + if (bytes > kMaxBytes - memory_usage_) { + return Status::Invalid("hash map total allocation size overflow"); + } + std::shared_ptr<Bytes> data = AllocateBytesKeepingPoolAlive(bytes, pool_); + auto* bitmap = reinterpret_cast<uint8_t*>(data->data() + bucket_count * kEntrySize); + std::memset(bitmap, 0, bitmap_size); + blocks_.push_back(Block{std::move(data), bitmap, bucket_count, 0}); + memory_usage_ += bytes; + if (initial_bucket_count_ == 0) { + initial_bucket_count_ = bucket_count; + } + return Status::OK(); + } + + std::shared_ptr<MemoryPool> pool_; + std::vector<Block> blocks_; + uint64_t initial_bucket_count_ = 0; + uint64_t size_ = 0; + uint64_t memory_usage_ = 0; +}; + +} // namespace paimon
diff --git a/src/paimon/common/utils/hash_map_test.cpp b/src/paimon/common/utils/hash_map_test.cpp new file mode 100644 index 0000000..b970d0d --- /dev/null +++ b/src/paimon/common/utils/hash_map_test.cpp
@@ -0,0 +1,132 @@ +/* + * 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/common/utils/hash_map.h" + +#include <cstdint> +#include <limits> +#include <map> +#include <memory> +#include <random> + +#include "gtest/gtest.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +using Map = HashMap<uint64_t, int64_t>; + +} // namespace + +TEST(HashMapTest, TestGrowthUpdatesAndIteration) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Map> map, Map::Create(1, GetDefaultPool())); + ASSERT_EQ(map->Size(), 0); + ASSERT_FALSE(map->Find(0)); + const uint64_t initial_memory = map->MemoryUsage(); + std::map<uint64_t, int64_t> expected; + for (uint64_t key = 0; key < 300; ++key) { + ASSERT_OK(map->FindAndInsert(key, key + 10)); + expected[key] = key + 10; + } + ASSERT_GT(map->MemoryUsage(), initial_memory); + const uint64_t memory_after_growth = map->MemoryUsage(); + for (uint64_t key = 0; key < 300; ++key) { + const int64_t value = -static_cast<int64_t>(key) - 1; + ASSERT_OK(map->FindAndInsert(key, value)); + expected[key] = value; + ASSERT_EQ(map->Find(key), value); + } + ASSERT_EQ(map->Size(), expected.size()); + ASSERT_EQ(map->MemoryUsage(), memory_after_growth); + ASSERT_FALSE(map->Find(300)); + + std::map<uint64_t, int64_t> visited; + map->ForEach( + [&](uint64_t key, int64_t value) { ASSERT_TRUE(visited.emplace(key, value).second); }); + ASSERT_EQ(visited, expected); +} + +TEST(HashMapTest, TestBucketCollisionsAndKeyBitPatterns) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Map> map, Map::Create(1, GetDefaultPool())); + // Distinct hashes with identical low bits start in the same bucket in these blocks. + for (uint64_t key = 0; key < 200; ++key) { + ASSERT_OK(map->FindAndInsert(key << 32, key * 2)); + } + const uint64_t max_key = std::numeric_limits<uint64_t>::max(); + ASSERT_OK(map->FindAndInsert(max_key, std::numeric_limits<int64_t>::min())); + for (uint64_t key = 0; key < 200; ++key) { + ASSERT_EQ(map->Find(key << 32), static_cast<int64_t>(key * 2)); + } + ASSERT_EQ(map->Find(max_key), std::numeric_limits<int64_t>::min()); + ASSERT_OK(map->FindAndInsert(0, std::numeric_limits<int64_t>::max())); + ASSERT_EQ(map->Find(0), std::numeric_limits<int64_t>::max()); + ASSERT_EQ(map->Size(), 201); + ASSERT_FALSE(map->Find(uint64_t{201} << 32)); +} + +TEST(HashMapTest, TestDeterministicUpdatesMatchReference) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Map> map, Map::Create(4, GetDefaultPool())); + std::map<uint64_t, int64_t> expected; + std::mt19937_64 random(42); + for (int32_t iteration = 0; iteration < 5000; ++iteration) { + const uint64_t key = random() % 1000; + const int64_t value = iteration - 2500; + ASSERT_OK(map->FindAndInsert(key, value)); + expected[key] = value; + } + ASSERT_EQ(map->Size(), expected.size()); + for (const auto& [key, value] : expected) { + ASSERT_EQ(map->Find(key), value); + } +} + +TEST(HashMapTest, TestClearAndPoolLifetime) { + std::shared_ptr<MemoryPool> pool = GetMemoryPool(); + std::weak_ptr<MemoryPool> weak_pool = pool; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<Map> map, Map::Create(2, pool)); + for (uint64_t key = 0; key < 100; ++key) { + ASSERT_OK(map->FindAndInsert(key, key)); + } + ASSERT_EQ(map->MemoryUsage(), pool->CurrentUsage()); + map->Clear(); + ASSERT_EQ(map->Size(), 0); + ASSERT_EQ(map->MemoryUsage(), 0); + ASSERT_EQ(pool->CurrentUsage(), 0); + ASSERT_FALSE(map->Find(1)); + uint64_t visited = 0; + map->ForEach([&](uint64_t, int64_t) { ++visited; }); + ASSERT_EQ(visited, 0); + ASSERT_OK(map->FindAndInsert(7, 9)); + ASSERT_EQ(map->Find(7), 9); + pool.reset(); + ASSERT_FALSE(weak_pool.expired()); + map.reset(); + ASSERT_TRUE(weak_pool.expired()); +} + +TEST(HashMapTest, TestInvalidArguments) { + ASSERT_NOK_WITH_MSG(Map::Create(1, nullptr), "requires a pool"); + ASSERT_NOK_WITH_MSG(Map::Create(0, GetDefaultPool()), "positive capacity"); + ASSERT_NOK(Map::Create(std::numeric_limits<uint64_t>::max(), GetDefaultPool())); + ASSERT_NOK(Map::Create(std::numeric_limits<int64_t>::max() / 2, GetDefaultPool())); +} + +} // namespace paimon::test
diff --git a/src/paimon/core/io/append_data_file_writer_factory.cpp b/src/paimon/core/io/append_data_file_writer_factory.cpp index d0b677a..e6abfa0 100644 --- a/src/paimon/core/io/append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/append_data_file_writer_factory.cpp
@@ -63,6 +63,11 @@ if (file_index_writer) { writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<DataFileKeyOffsetIndexWriter> offset_index_writer, + CreateOffsetIndexWriter(write_schema_)); + if (offset_index_writer) { + writer->SetOffsetIndexWriter(std::move(offset_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); return std::unique_ptr<SingleFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>(
diff --git a/src/paimon/core/io/data_file_writer.cpp b/src/paimon/core/io/data_file_writer.cpp index 5a15176..2263f0e 100644 --- a/src/paimon/core/io/data_file_writer.cpp +++ b/src/paimon/core/io/data_file_writer.cpp
@@ -67,7 +67,7 @@ return DataFileMeta::ForAppend( PathUtil::GetName(path_), output_bytes_, RecordCount(), stats, seq_num_counter_->GetValue() - RecordCount(), seq_num_counter_->GetValue() - 1, schema_id_, - file_index.extra_files, file_index.embedded_index, file_source_, + GetExtraFiles(), file_index.embedded_index, file_source_, /*value_stats_cols=*/std::nullopt, final_path, /*first_row_id=*/std::nullopt, write_cols_); }
diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h index 3583c63..6eb8eaa 100644 --- a/src/paimon/core/io/data_file_writer_base.h +++ b/src/paimon/core/io/data_file_writer_base.h
@@ -20,9 +20,11 @@ #include <functional> #include <memory> +#include <optional> #include <string> #include <type_traits> #include <utility> +#include <vector> #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" @@ -32,6 +34,7 @@ #include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/single_file_writer.h" +#include "paimon/core/realtime/data_file_key_offset_index_writer.h" #include "paimon/format/format_stats_extractor.h" #include "paimon/format/written_file_stats_provider.h" #include "paimon/result.h" @@ -42,8 +45,8 @@ struct KeyValueBatch; -/// Common lifecycle for data file writers which may finalize schema metadata and publish a -/// file-level index. Concrete writers remain responsible for their record-specific state and +/// Common lifecycle for data file writers which may finalize schema metadata and publish +/// file-level indexes. Concrete writers remain responsible for their record-specific state and /// DataFileMeta construction. template <typename Record> class DataFileWriterBase : public SingleFileWriter<Record, std::shared_ptr<DataFileMeta>> { @@ -68,6 +71,9 @@ } void Abort() override { + if (offset_index_writer_) { + offset_index_writer_->Abort(); + } if (file_index_writer_) { // The external index uses a path different from the data file path deleted by Base. file_index_writer_->Abort(); @@ -80,15 +86,24 @@ if (file_index_writer_ && file_index_writer_->ExternalIndexPath()) { executor.Add(this->fs_, file_index_writer_->ExternalIndexPath().value()); } + if (offset_index_writer_ && !offset_index_writer_->OutputPath().empty()) { + executor.Add(this->fs_, offset_index_writer_->OutputPath()); + } return executor; } + void SetOffsetIndexWriter(std::unique_ptr<DataFileKeyOffsetIndexWriter>&& offset_index_writer, + const std::shared_ptr<arrow::Schema>& logical_schema) { + offset_index_writer_ = std::move(offset_index_writer); + logical_type_ = arrow::struct_(logical_schema->fields()); + } + protected: DataFileWriterBase(const std::string& compression, std::function<Status(Record, ::ArrowArray*)> converter) : Base(compression, std::move(converter)) {} - /// Extracts the pre-conversion Arrow batch from record for file index construction, then + /// Extracts the pre-conversion Arrow batch for file and offset index construction, then /// passes record to the underlying data file writer, which may convert it to a physical schema. Status WriteRecordWithFileIndex(Record record) { PAIMON_RETURN_NOT_OK(this->CheckNotClosed()); @@ -100,6 +115,14 @@ return file_index_result_; } + std::vector<std::optional<std::string>> GetExtraFiles() const { + std::vector<std::optional<std::string>> extra_files = file_index_result_.extra_files; + if (offset_index_file_) { + extra_files.push_back(offset_index_file_); + } + return extra_files; + } + /// Extracts the column statistics of the closed file from the format writer when it provides /// them, and reads them back from the file through `stats_extractor` otherwise. Result<ColumnStatsVector> ExtractFileStats( @@ -128,6 +151,10 @@ if (file_index_writer_) { PAIMON_ASSIGN_OR_RAISE(file_index_result_, file_index_writer_->Finish(this->path_)); } + if (offset_index_writer_) { + PAIMON_ASSIGN_OR_RAISE(offset_index_file_, + offset_index_writer_->Finish(this->path_, this->RecordCount())); + } return Status::OK(); } @@ -143,7 +170,7 @@ } Status AddFileIndexBatch(::ArrowArray* batch) { - if (!file_index_writer_) { + if (!file_index_writer_ && !offset_index_writer_) { return Status::OK(); } // A compaction rewrite forwards the dictionary encoding of its input files, which the @@ -162,13 +189,20 @@ arrow::ImportArray(batch, batch_type.value())); std::shared_ptr<arrow::StructArray> logical_batch = checked_pointer_cast<arrow::StructArray>(logical_array); - PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); + if (file_index_writer_) { + PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); + } + if (offset_index_writer_) { + PAIMON_RETURN_NOT_OK(offset_index_writer_->AddBatch(logical_batch)); + } PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*logical_batch, batch)); return Status::OK(); } MetadataFinalizer metadata_finalizer_; std::unique_ptr<DataFileIndexWriter> file_index_writer_; + std::unique_ptr<DataFileKeyOffsetIndexWriter> offset_index_writer_; + std::optional<std::string> offset_index_file_; std::shared_ptr<arrow::DataType> logical_type_; FileIndexWriteResult file_index_result_; };
diff --git a/src/paimon/core/io/data_file_writer_base_test.cpp b/src/paimon/core/io/data_file_writer_base_test.cpp index f2b5ac3..dffa3cc 100644 --- a/src/paimon/core/io/data_file_writer_base_test.cpp +++ b/src/paimon/core/io/data_file_writer_base_test.cpp
@@ -19,10 +19,13 @@ #include "paimon/core/io/data_file_writer_base.h" #include <cstdint> +#include <limits> +#include <map> #include <memory> #include <optional> #include <string> #include <utility> +#include <vector> #include "arrow/api.h" #include "arrow/c/bridge.h" @@ -30,6 +33,7 @@ #include "gtest/gtest.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/long_counter.h" #include "paimon/core/core_options.h" #include "paimon/core/io/append_data_file_writer_factory.h" @@ -37,8 +41,10 @@ #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/file_index_options.h" #include "paimon/core/io/key_value_data_file_writer_factory.h" +#include "paimon/core/io/rolling_file_writer.h" #include "paimon/core/key_value.h" #include "paimon/core/manifest/file_source.h" +#include "paimon/core/realtime/realtime_key_offset_index.h" #include "paimon/core/stats/simple_stats.h" #include "paimon/core/stats/simple_stats_converter.h" #include "paimon/defs.h" @@ -65,6 +71,20 @@ mutable int32_t open_count = 0; }; +class FailingOffsetFileSystem : public LocalFileSystem { + public: + Result<std::unique_ptr<OutputStream>> Create(const std::string& path, + bool overwrite) const override { + if (path.size() >= 7 && path.compare(path.size() - 7, 7, ".offset") == 0 && + ++offset_files == 2) { + return Status::IOError("injected offset index create failure"); + } + return LocalFileSystem::Create(path, overwrite); + } + + mutable int32_t offset_files = 0; +}; + } // namespace class IndexedDataFileWriter : public DataFileWriterBase<::ArrowArray*> { @@ -125,6 +145,117 @@ ArrowArrayRelease(&second_batch); } +TEST(DataFileWriterBaseTest, OffsetIndexesFollowRollingWithAndWithoutFileIndexes) { + for (const std::string index_storage : {"none", "embedded", "external"}) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + const std::shared_ptr<MemoryPool> pool = GetDefaultPool(); + const std::shared_ptr<arrow::Schema> schema = + arrow::schema({arrow::field(SpecialFields::RealtimeOffset().Name(), arrow::int64()), + arrow::field("id", arrow::utf8())}); + std::map<std::string, std::string> raw_options = { + {Options::FILE_FORMAT, "parquet"}, {Options::REALTIME_DEDUPLICATE_KEY_FIELDS, "id"}}; + if (index_storage != "none") { + raw_options["file-index.bitmap.columns"] = "id"; + raw_options[Options::FILE_INDEX_IN_MANIFEST_THRESHOLD] = + index_storage == "embedded" ? "1MB" : "0"; + } + ASSERT_OK_AND_ASSIGN(CoreOptions options, CoreOptions::FromMap(raw_options)); + auto path_factory = std::make_shared<DataFilePathFactory>(); + ASSERT_OK(path_factory->Init(dir->Str(), "parquet", "data-", nullptr)); + auto factory = std::make_shared<AppendDataFileWriterFactory>( + options, /*schema_id=*/0, schema, /*write_cols=*/std::nullopt, + std::make_shared<LongCounter>(0), FileSource::Append(), path_factory, pool); + RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>> writer( + std::numeric_limits<int64_t>::max(), /*target_file_row_num=*/2, factory); + for (const std::string json : + {R"([[100, "a"], [102, "b"]])", R"([[200, "c"], [203, "d"]])"}) { + const std::shared_ptr<arrow::Array> array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), json) + .ValueOrDie(); + ::ArrowArray batch; + ASSERT_TRUE(arrow::ExportArray(*array, &batch).ok()); + ASSERT_OK(writer.Write(&batch)); + } + ASSERT_OK(writer.Close()); + ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<DataFileMeta>> files, writer.GetResult()); + ASSERT_EQ(2, files.size()); + auto cache_manager = std::make_shared<CacheManager>(1024, 0.0); + const std::shared_ptr<arrow::StructArray> keys = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({schema->field(1)}), + R"([["a"], ["b"], ["c"], ["d"]])") + .ValueOrDie()); + const std::vector<std::vector<int64_t>> expected_offsets = {{100, 102}, {200, 203}}; + for (size_t index = 0; index < files.size(); ++index) { + const std::shared_ptr<DataFileMeta>& file = files[index]; + ASSERT_EQ(2, file->row_count); + ASSERT_EQ(index_storage == "external" ? 2 : 1, file->extra_files.size()); + ASSERT_EQ(index_storage == "embedded", file->embedded_index != nullptr); + ASSERT_EQ(file->file_name + ".offset", file->extra_files.back().value()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create(32, cache_manager, schema->field(1), file, + path_factory, options.GetFileSystem(), pool)); + ASSERT_OK_AND_ASSIGN(RoaringBitmap64 offsets, reader->LookupOffsets(keys)); + ASSERT_EQ(2, offsets.Cardinality()); + for (int64_t offset : expected_offsets[index]) { + ASSERT_TRUE(offsets.Contains(offset)); + } + for (const std::string& path : path_factory->CollectFiles(file)) { + ASSERT_OK_AND_ASSIGN(bool exists, options.GetFileSystem()->Exists(path)); + ASSERT_TRUE(exists); + } + } + // A rolling writer must retain sidecar paths in the closed writers' abort executors. + writer.Abort(); + for (const std::shared_ptr<DataFileMeta>& file : files) { + for (const std::string& path : path_factory->CollectFiles(file)) { + ASSERT_OK_AND_ASSIGN(bool exists, options.GetFileSystem()->Exists(path)); + ASSERT_FALSE(exists); + } + } + } +} + +TEST(DataFileWriterBaseTest, OffsetIndexFailureCleansCurrentAndPreviouslyRolledFiles) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto fs = std::make_shared<FailingOffsetFileSystem>(); + const std::shared_ptr<arrow::Schema> schema = + arrow::schema({arrow::field(SpecialFields::RealtimeOffset().Name(), arrow::int64()), + arrow::field("id", arrow::int64())}); + ASSERT_OK_AND_ASSIGN(CoreOptions options, + CoreOptions::FromMap({{Options::FILE_FORMAT, "parquet"}, + {Options::REALTIME_DEDUPLICATE_KEY_FIELDS, "id"}, + {"file-index.bitmap.columns", "id"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "0"}}, + fs)); + auto path_factory = std::make_shared<DataFilePathFactory>(); + ASSERT_OK(path_factory->Init(dir->Str(), "parquet", "data-", nullptr)); + auto factory = std::make_shared<AppendDataFileWriterFactory>( + options, /*schema_id=*/0, schema, /*write_cols=*/std::nullopt, + std::make_shared<LongCounter>(0), FileSource::Append(), path_factory, GetDefaultPool()); + RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>> writer( + std::numeric_limits<int64_t>::max(), /*target_file_row_num=*/1, factory); + for (int32_t index = 0; index < 2; ++index) { + const std::shared_ptr<arrow::Array> array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(schema->fields()), + R"([[7, 1]])") + .ValueOrDie(); + ::ArrowArray batch; + ASSERT_TRUE(arrow::ExportArray(*array, &batch).ok()); + if (index == 0) { + ASSERT_OK(writer.Write(&batch)); + } else { + ASSERT_NOK_WITH_MSG(writer.Write(&batch), "injected offset index create failure"); + } + } + ASSERT_EQ(2, fs->offset_files); + std::vector<BasicFileStatus> files; + ASSERT_OK(fs->ListDir(dir->Str(), &files)); + ASSERT_TRUE(files.empty()); +} + class DataFileWriterStatsTest : public ::testing::TestWithParam<std::string> { public: void SetUp() override {
diff --git a/src/paimon/core/io/data_file_writer_factory.cpp b/src/paimon/core/io/data_file_writer_factory.cpp index 07195ab..aa273d1 100644 --- a/src/paimon/core/io/data_file_writer_factory.cpp +++ b/src/paimon/core/io/data_file_writer_factory.cpp
@@ -21,12 +21,14 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "arrow/type.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/scope_guard.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/io/file_index_options.h" +#include "paimon/core/realtime/data_file_key_offset_index_writer.h" #include "paimon/format/file_format.h" #include "paimon/format/writer_builder.h" @@ -73,4 +75,18 @@ path_factory, pool_); } +Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> +DataFileWriterFactory::CreateOffsetIndexWriter( + const std::shared_ptr<arrow::Schema>& logical_schema) const { + const std::vector<std::string>& key_fields = options_.GetRealtimeDeduplicateKeyFields(); + if (key_fields.empty()) { + return std::unique_ptr<DataFileKeyOffsetIndexWriter>(); + } + if (!logical_schema || key_fields.size() != 1) { + return Status::Invalid("offset index requires exactly one deduplicate key field"); + } + std::shared_ptr<arrow::Field> key_field = logical_schema->GetFieldByName(key_fields.front()); + return DataFileKeyOffsetIndexWriter::Create(key_field, options_.GetFileSystem(), pool_); +} + } // namespace paimon
diff --git a/src/paimon/core/io/data_file_writer_factory.h b/src/paimon/core/io/data_file_writer_factory.h index cab942f..94b0ed9 100644 --- a/src/paimon/core/io/data_file_writer_factory.h +++ b/src/paimon/core/io/data_file_writer_factory.h
@@ -33,6 +33,7 @@ class FileFormat; class DataFileIndexWriter; +class DataFileKeyOffsetIndexWriter; class DataFilePathFactory; class FormatStatsExtractor; class MemoryPool; @@ -58,6 +59,9 @@ const std::shared_ptr<arrow::Schema>& logical_schema, const std::shared_ptr<DataFilePathFactory>& path_factory) const; + Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> CreateOffsetIndexWriter( + const std::shared_ptr<arrow::Schema>& logical_schema) const; + CoreOptions options_; int64_t schema_id_; std::shared_ptr<MemoryPool> pool_;
diff --git a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp index 0e4e821..43f2c37 100644 --- a/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp +++ b/src/paimon/core/io/shredding_append_data_file_writer_factory.cpp
@@ -95,6 +95,11 @@ if (file_index_writer) { writer->SetFileIndexWriter(std::move(file_index_writer), write_schema_); } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<DataFileKeyOffsetIndexWriter> offset_index_writer, + CreateOffsetIndexWriter(write_schema_)); + if (offset_index_writer) { + writer->SetOffsetIndexWriter(std::move(offset_index_writer), write_schema_); + } PAIMON_RETURN_NOT_OK( writer->Init(options_.GetFileSystem(), path_factory_->NewPath(), resources.writer_builder)); ShreddingWritePlanFactory::MetadataFinalizer finalizer =
diff --git a/src/paimon/core/realtime/data_file_key_offset_index_writer.cpp b/src/paimon/core/realtime/data_file_key_offset_index_writer.cpp new file mode 100644 index 0000000..85d83a6 --- /dev/null +++ b/src/paimon/core/realtime/data_file_key_offset_index_writer.cpp
@@ -0,0 +1,170 @@ +/* + * 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/data_file_key_offset_index_writer.h" + +#include <utility> + +#include "arrow/api.h" +#include "arrow/compute/cast.h" +#include "paimon/common/global_index/key_serializer.h" +#include "paimon/common/io/data_output_stream.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/dense_hash_table.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/casting/casting_utils.h" +#include "paimon/core/realtime/key_offset_index_utils.h" +#include "paimon/fs/file_system.h" +#include "paimon/macros.h" + +namespace paimon { +Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> DataFileKeyOffsetIndexWriter::Create( + const std::shared_ptr<arrow::Field>& key_field, const std::shared_ptr<FileSystem>& file_system, + const std::shared_ptr<MemoryPool>& memory_pool) { + if (!key_field || !file_system || !memory_pool) { + return Status::Invalid("key-offset writer is missing a required dependency"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeySerializer> serializer, + KeySerializer::Create(key_field->type(), memory_pool)); + using OffsetMap = HashMap<uint64_t, int64_t>; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<OffsetMap> offsets, + OffsetMap::Create(/*initial_capacity=*/1024, memory_pool)); + return std::unique_ptr<DataFileKeyOffsetIndexWriter>(new DataFileKeyOffsetIndexWriter( + key_field, file_system, serializer, std::move(offsets), memory_pool)); +} + +DataFileKeyOffsetIndexWriter::DataFileKeyOffsetIndexWriter( + const std::shared_ptr<arrow::Field>& key_field, const std::shared_ptr<FileSystem>& file_system, + const std::shared_ptr<KeySerializer>& key_serializer, + std::unique_ptr<HashMap<uint64_t, int64_t>> offsets, + const std::shared_ptr<MemoryPool>& memory_pool) + : key_field_(key_field), + file_system_(file_system), + key_serializer_(key_serializer), + offsets_(std::move(offsets)), + memory_pool_(memory_pool) {} + +Status DataFileKeyOffsetIndexWriter::AddBatch( + const std::shared_ptr<arrow::StructArray>& logical_batch) { + if (finished_) { + return Status::Invalid("key-offset writer has already finished"); + } + if (!logical_batch || logical_batch->null_count() != 0) { + return Status::Invalid("key-offset writer requires non-null logical rows"); + } + // The adaptor must outlive any decoded arrays allocated through it. + std::shared_ptr<arrow::MemoryPool> arrow_pool; + std::shared_ptr<arrow::Array> key = logical_batch->GetFieldByName(key_field_->name()); + std::shared_ptr<arrow::Array> offsets = + logical_batch->GetFieldByName(SpecialFields::RealtimeOffset().Name()); + if (!key || !offsets) { + return Status::Invalid("key-offset writer requires the key and realtime offset columns"); + } + // Compaction may forward dictionary-encoded columns. Materialize only the index inputs. + if (key->type_id() == arrow::Type::DICTIONARY || + offsets->type_id() == arrow::Type::DICTIONARY) { + arrow_pool = GetArrowPool(memory_pool_); + } + if (key->type_id() == arrow::Type::DICTIONARY) { + PAIMON_ASSIGN_OR_RAISE( + key, CastingUtils::Cast(key, key_field_->type(), arrow::compute::CastOptions::Safe(), + arrow_pool.get())); + } + if (offsets->type_id() == arrow::Type::DICTIONARY) { + PAIMON_ASSIGN_OR_RAISE( + offsets, CastingUtils::Cast(offsets, arrow::int64(), + arrow::compute::CastOptions::Safe(), arrow_pool.get())); + } + if (offsets->type_id() != arrow::Type::INT64) { + return Status::Invalid("realtime offset must have int64 type"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> keys, + arrow::StructArray::Make({key}, {key_field_})); + const std::shared_ptr<arrow::Int64Array> typed_offsets = + checked_pointer_cast<arrow::Int64Array>(offsets); + if (keys->length() != typed_offsets->length() || typed_offsets->null_count() != 0) { + return Status::Invalid("key-offset writer input columns are not aligned"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector<uint64_t> hashes, + KeyOffsetIndexUtils::HashKeys(keys, key_field_, key_serializer_)); + for (int64_t row = 0; row < typed_offsets->length(); ++row) { + const int64_t offset = typed_offsets->Value(row); + if (offset < 0) { + return Status::Invalid("realtime offset must not be negative"); + } + if (offsets_->Find(hashes[row])) { + return Status::Invalid( + "deduplicate data file contains more than one offset for a deduplicate key"); + } + PAIMON_RETURN_NOT_OK(offsets_->FindAndInsert(hashes[row], offset)); + } + return Status::OK(); +} + +Result<std::string> DataFileKeyOffsetIndexWriter::Finish(const std::string& data_file_path, + int64_t row_count) { + if (finished_) { + return Status::Invalid("key-offset writer has already finished"); + } + if (data_file_path.empty() || offsets_->Size() == 0) { + return Status::Invalid("cannot finish an empty key-offset index"); + } + if (row_count < 0 || offsets_->Size() != static_cast<uint64_t>(row_count)) { + return Status::Invalid("realtime .offset index entry count does not match data file"); + } + finished_ = true; + using HashTable = DenseHashTable<uint64_t, int64_t>; + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<HashTable> hash_table, + HashTable::Create(offsets_->Size(), /*occupancy_pct=*/80, memory_pool_)); + Status status = Status::OK(); + offsets_->ForEach([&](uint64_t hash, int64_t offset) { + if (status.ok()) { + status = hash_table->Insert(hash, offset); + } + }); + PAIMON_RETURN_NOT_OK(status); + const std::string output_path = data_file_path + KeyOffsetIndexUtils::kFileSuffix; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<OutputStream> output, + file_system_->Create(output_path, /*overwrite=*/false)); + output_path_ = output_path; + ScopeGuard guard([this, &output]() { + [[maybe_unused]] Status status = output->Close(); + Abort(); + }); + DataOutputStream data_output(output); + PAIMON_RETURN_NOT_OK(data_output.WriteValue<int32_t>(KeyOffsetIndexUtils::kCurrentVersion)); + PAIMON_RETURN_NOT_OK(hash_table->WriteTo(output.get())); + PAIMON_RETURN_NOT_OK(output->Flush()); + PAIMON_RETURN_NOT_OK(output->Close()); + guard.Release(); + return PathUtil::GetName(output_path_); +} + +void DataFileKeyOffsetIndexWriter::Abort() { + if (!output_path_.empty()) { + [[maybe_unused]] Status status = file_system_->Delete(output_path_); + output_path_.clear(); + } +} + +} // namespace paimon
diff --git a/src/paimon/core/realtime/data_file_key_offset_index_writer.h b/src/paimon/core/realtime/data_file_key_offset_index_writer.h new file mode 100644 index 0000000..2c2aded --- /dev/null +++ b/src/paimon/core/realtime/data_file_key_offset_index_writer.h
@@ -0,0 +1,68 @@ +/* + * 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 <cstdint> +#include <memory> +#include <string> + +#include "arrow/type_fwd.h" +#include "paimon/common/utils/hash_map.h" +#include "paimon/result.h" + +namespace paimon { +class FileSystem; +class KeySerializer; +class MemoryPool; + +/// Builds one data file's hash -> global realtime offset index from its surviving logical rows. +/// Layout: [big-endian int32 version][native DenseHashTable<uint64_t, int64_t> layout]. +class DataFileKeyOffsetIndexWriter { + public: + static Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> Create( + const std::shared_ptr<arrow::Field>& key_field, + const std::shared_ptr<FileSystem>& file_system, + const std::shared_ptr<MemoryPool>& memory_pool); + + Status AddBatch(const std::shared_ptr<arrow::StructArray>& logical_batch); + + Result<std::string> Finish(const std::string& data_file_path, int64_t row_count); + void Abort(); + + const std::string& OutputPath() const { + return output_path_; + } + + private: + DataFileKeyOffsetIndexWriter(const std::shared_ptr<arrow::Field>& key_field, + const std::shared_ptr<FileSystem>& file_system, + const std::shared_ptr<KeySerializer>& key_serializer, + std::unique_ptr<HashMap<uint64_t, int64_t>> offsets, + const std::shared_ptr<MemoryPool>& memory_pool); + + std::shared_ptr<arrow::Field> key_field_; + std::shared_ptr<FileSystem> file_system_; + std::shared_ptr<KeySerializer> key_serializer_; + std::unique_ptr<HashMap<uint64_t, int64_t>> offsets_; + std::shared_ptr<MemoryPool> memory_pool_; + std::string output_path_; + bool finished_ = false; +}; +} // namespace paimon
diff --git a/src/paimon/core/realtime/key_offset_index_utils.cpp b/src/paimon/core/realtime/key_offset_index_utils.cpp new file mode 100644 index 0000000..bd02cf7 --- /dev/null +++ b/src/paimon/core/realtime/key_offset_index_utils.cpp
@@ -0,0 +1,51 @@ +/* + * 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/key_offset_index_utils.h" + +#include "arrow/api.h" +#include "paimon/common/global_index/key_serializer.h" +#include "paimon/common/predicate/literal_converter.h" +#include "paimon/macros.h" +#include "xxhash.h" // NOLINT(build/include_subdir) + +namespace paimon { +Result<std::vector<uint64_t>> KeyOffsetIndexUtils::HashKeys( + const std::shared_ptr<arrow::StructArray>& keys, const std::shared_ptr<arrow::Field>& key_field, + const std::shared_ptr<KeySerializer>& key_serializer) { + if (!keys || keys->num_fields() != 1 || !keys->field(0)->type()->Equals(key_field->type())) { + return Status::Invalid( + "deduplicate key lookup must contain exactly the configured key field"); + } + if (keys->null_count() != 0 || keys->field(0)->null_count() != 0) { + return Status::Invalid("deduplicate key must not contain null"); + } + PAIMON_ASSIGN_OR_RAISE(std::vector<Literal> literals, + LiteralConverter::ConvertLiteralsFromArray(*keys->field(0), + /*own_data=*/true)); + std::vector<uint64_t> hashes; + hashes.reserve(literals.size()); + for (const Literal& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Bytes> bytes, key_serializer->Serialize(literal)); + hashes.push_back(XXH64(bytes->data(), bytes->size(), /*seed=*/0)); + } + return hashes; +} + +} // namespace paimon
diff --git a/src/paimon/core/realtime/key_offset_index_utils.h b/src/paimon/core/realtime/key_offset_index_utils.h new file mode 100644 index 0000000..ee72cd8 --- /dev/null +++ b/src/paimon/core/realtime/key_offset_index_utils.h
@@ -0,0 +1,46 @@ +/* + * 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 <cstdint> +#include <memory> +#include <vector> + +#include "arrow/type_fwd.h" +#include "paimon/result.h" + +namespace paimon { +class KeySerializer; + +/// Shared hash algorithm and file format constants for realtime and per-file offset indexes. +class KeyOffsetIndexUtils { + public: + KeyOffsetIndexUtils() = delete; + ~KeyOffsetIndexUtils() = delete; + + static constexpr const char* kFileSuffix = ".offset"; + static constexpr int32_t kCurrentVersion = 1; + + static Result<std::vector<uint64_t>> HashKeys( + const std::shared_ptr<arrow::StructArray>& keys, + const std::shared_ptr<arrow::Field>& key_field, + const std::shared_ptr<KeySerializer>& key_serializer); +}; +} // namespace paimon
diff --git a/src/paimon/core/realtime/key_offset_index_utils_test.cpp b/src/paimon/core/realtime/key_offset_index_utils_test.cpp new file mode 100644 index 0000000..b45789c --- /dev/null +++ b/src/paimon/core/realtime/key_offset_index_utils_test.cpp
@@ -0,0 +1,109 @@ +/* + * 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/key_offset_index_utils.h" + +#include <cstdint> +#include <memory> +#include <string> +#include <vector> + +#include "arrow/api.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/global_index/key_serializer.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/testing/utils/testharness.h" +#include "xxhash.h" // NOLINT(build/include_subdir) + +namespace paimon::test { +namespace { + +std::shared_ptr<arrow::StructArray> MakeKeys(const std::shared_ptr<arrow::Field>& field, + const std::string& json) { + return checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({field}), json).ValueOrDie()); +} + +} // namespace + +TEST(KeyOffsetIndexUtilsTest, TestIntegerHashesMatchSerializedKeysAndPreserveOrder) { + const std::shared_ptr<arrow::Field> field = arrow::field("id", arrow::int64()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<KeySerializer> serializer, + KeySerializer::Create(field->type(), GetDefaultPool())); + const std::shared_ptr<arrow::StructArray> keys = MakeKeys(field, "[[42], [-7], [42], [0]]"); + ASSERT_OK_AND_ASSIGN(std::vector<uint64_t> hashes, + KeyOffsetIndexUtils::HashKeys(keys, field, serializer)); + std::vector<uint64_t> expected; + for (int64_t value : {42, -7, 42, 0}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> bytes, serializer->Serialize(Literal(value))); + expected.push_back(XXH64(bytes->data(), bytes->size(), /*seed=*/0)); + } + ASSERT_EQ(expected, hashes); + + const std::shared_ptr<arrow::StructArray> slice = + checked_pointer_cast<arrow::StructArray>(keys->Slice(1, 2)); + ASSERT_OK_AND_ASSIGN(std::vector<uint64_t> sliced_hashes, + KeyOffsetIndexUtils::HashKeys(slice, field, serializer)); + const std::vector<uint64_t> expected_slice = {expected[1], expected[2]}; + ASSERT_EQ(expected_slice, sliced_hashes); +} + +TEST(KeyOffsetIndexUtilsTest, TestStringKeysAndEmptyBatch) { + const std::shared_ptr<arrow::Field> field = arrow::field("id", arrow::utf8()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<KeySerializer> serializer, + KeySerializer::Create(field->type(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::vector<uint64_t> hashes, + KeyOffsetIndexUtils::HashKeys(MakeKeys(field, R"([[""], ["alpha"], ["alpha"]])"), field, + serializer)); + const std::vector<uint64_t> expected = {XXH64("", 0, /*seed=*/0), XXH64("alpha", 5, /*seed=*/0), + XXH64("alpha", 5, /*seed=*/0)}; + ASSERT_EQ(expected, hashes); + ASSERT_OK_AND_ASSIGN(hashes, + KeyOffsetIndexUtils::HashKeys(MakeKeys(field, "[]"), field, serializer)); + ASSERT_TRUE(hashes.empty()); +} + +TEST(KeyOffsetIndexUtilsTest, TestRejectsInvalidKeyBatches) { + const std::shared_ptr<arrow::Field> field = arrow::field("id", arrow::int64()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<KeySerializer> serializer, + KeySerializer::Create(field->type(), GetDefaultPool())); + ASSERT_NOK_WITH_MSG(KeyOffsetIndexUtils::HashKeys(nullptr, field, serializer), + "exactly the configured key field"); + const std::shared_ptr<arrow::StructArray> multiple_fields = + checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({field, arrow::field("other", arrow::int64())}), "[[1, 2]]") + .ValueOrDie()); + ASSERT_NOK_WITH_MSG(KeyOffsetIndexUtils::HashKeys(multiple_fields, field, serializer), + "exactly the configured key field"); + ASSERT_NOK_WITH_MSG( + KeyOffsetIndexUtils::HashKeys(MakeKeys(arrow::field("id", arrow::utf8()), R"([["1"]])"), + field, serializer), + "exactly the configured key field"); + ASSERT_NOK_WITH_MSG( + KeyOffsetIndexUtils::HashKeys(MakeKeys(field, "[[null]]"), field, serializer), + "deduplicate key must not contain null"); + ASSERT_NOK_WITH_MSG(KeyOffsetIndexUtils::HashKeys(MakeKeys(field, "[null]"), field, serializer), + "deduplicate key must not contain null"); +} + +} // namespace paimon::test
diff --git a/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp b/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp index 45dde7b..d15a179 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp +++ b/src/paimon/core/realtime/realtime_deduplicate_state_test.cpp
@@ -34,6 +34,7 @@ #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/realtime/arrow_deduplicate_realtime_store.h" #include "paimon/core/realtime/arrow_realtime_store.h" +#include "paimon/core/realtime/data_file_key_offset_index_writer.h" #include "paimon/core/realtime/realtime_key_offset_index.h" #include "paimon/core/realtime/realtime_offset_file_index_lookup.h" #include "paimon/memory/memory_pool.h" @@ -97,15 +98,19 @@ /*write_cols=*/std::nullopt)); ASSERT_OK_AND_ASSIGN( std::unique_ptr<DataFileKeyOffsetIndexWriter> sidecar_writer, - DataFileKeyOffsetIndexWriter::Create(schema->field(1), path_factory, - directory->GetFileSystem(), pool, options.ToMap())); + DataFileKeyOffsetIndexWriter::Create(schema->field(1), directory->GetFileSystem(), pool)); const std::shared_ptr<arrow::StructArray> keys = checked_pointer_cast<arrow::StructArray>( arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({schema->field(1)}), R"([[1]])") .ValueOrDie()); const std::shared_ptr<arrow::Int64Array> offsets = checked_pointer_cast<arrow::Int64Array>( arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[0]").ValueOrDie()); - ASSERT_OK(sidecar_writer->AddBatch(keys, offsets)); - ASSERT_OK_AND_ASSIGN(std::string sidecar, sidecar_writer->Finish(first_file)); + const std::shared_ptr<arrow::StructArray> index_batch = + arrow::StructArray::Make({keys->field(0), offsets}, {schema->field(1), schema->field(0)}) + .ValueOrDie(); + ASSERT_OK(sidecar_writer->AddBatch(index_batch)); + ASSERT_OK_AND_ASSIGN( + std::string sidecar, + sidecar_writer->Finish(path_factory->ToPath(first_file), first_file->row_count)); first_file = first_file->CopyWithExtraFiles({sidecar}); ASSERT_NOK_WITH_MSG(state.InstallCommittedSnapshot(-1, {}, store), @@ -120,6 +125,8 @@ ASSERT_OK(state.InstallCommittedSnapshot(/*committed_end_offset=*/1, {first_file}, store)); const std::shared_ptr<RealtimeOffsetFileIndexLookup> installed = state.committed_file_lookup_; ASSERT_NE(initial_lookup, installed); + ASSERT_EQ(initial_lookup->cache_manager_, installed->cache_manager_); + ASSERT_EQ(options.GetCachePageSize(), installed->cache_page_size_); ASSERT_EQ(std::vector<std::shared_ptr<DataFileMeta>>{first_file}, installed->DataFiles()); ASSERT_OK_AND_ASSIGN(std::shared_ptr<const KeyOffsetLookup> new_keys, store->AcquireKeyOffsetLookup()); @@ -131,6 +138,7 @@ // A new snapshot may replace its active file set without advancing the durable offset. ASSERT_OK(state.InstallCommittedSnapshot(/*committed_end_offset=*/1, {}, store)); ASSERT_TRUE(state.committed_file_lookup_->DataFiles().empty()); + ASSERT_EQ(installed->cache_manager_, state.committed_file_lookup_->cache_manager_); ASSERT_EQ(std::vector<std::shared_ptr<DataFileMeta>>{first_file}, installed->DataFiles()); ASSERT_OK_AND_ASSIGN(RoaringBitmap64 pinned_offsets, new_keys->LookupOffsets(keys)); ASSERT_EQ("{0}", pinned_offsets.ToString());
diff --git a/src/paimon/core/realtime/realtime_deduplicate_writer.cpp b/src/paimon/core/realtime/realtime_deduplicate_writer.cpp index c908054..15eb933 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_writer.cpp +++ b/src/paimon/core/realtime/realtime_deduplicate_writer.cpp
@@ -25,7 +25,6 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "paimon/common/reader/concat_batch_reader.h" -#include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" @@ -109,27 +108,20 @@ store_state.deduplicate_state->AttachFileIndexLookup(created_lookup, store)); return std::shared_ptr<RealtimeDeduplicateWriter>(new RealtimeDeduplicateWriter( store, file_writer, schema_layout, schema_layout->StoreWriteSchema()->field(key_position), - data_file_path_factory, options.GetFileSystem(), store_state.initial_offset, memory_pool, - options.ToMap())); + store_state.initial_offset, memory_pool)); } RealtimeDeduplicateWriter::RealtimeDeduplicateWriter( std::shared_ptr<ArrowDeduplicateRealtimeStore> realtime_store, std::shared_ptr<AppendOnlyWriter> file_writer, std::shared_ptr<RealtimeSchemaLayout> schema_layout, - std::shared_ptr<arrow::Field> deduplicate_key_field, - std::shared_ptr<DataFilePathFactory> data_file_path_factory, - std::shared_ptr<FileSystem> file_system, int64_t next_offset, - std::shared_ptr<MemoryPool> memory_pool, std::map<std::string, std::string> options) + std::shared_ptr<arrow::Field> deduplicate_key_field, int64_t next_offset, + std::shared_ptr<MemoryPool> memory_pool) : arrow_pool_(GetArrowPool(memory_pool)), realtime_store_(std::move(realtime_store)), file_writer_(std::move(file_writer)), schema_layout_(std::move(schema_layout)), deduplicate_key_field_(std::move(deduplicate_key_field)), - data_file_path_factory_(std::move(data_file_path_factory)), - file_system_(std::move(file_system)), - memory_pool_(std::move(memory_pool)), - options_(std::move(options)), next_offset_(next_offset) {} Status RealtimeDeduplicateWriter::Write(std::unique_ptr<RecordBatch>&& batch) { @@ -221,41 +213,29 @@ if (segment->GetRowCount() == 0) { continue; } - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr<DataFileKeyOffsetIndexWriter> offset_index_writer, - DataFileKeyOffsetIndexWriter::Create(deduplicate_key_field_, data_file_path_factory_, - file_system_, memory_pool_, options_)); - ScopeGuard offset_index_guard([&offset_index_writer]() { - if (offset_index_writer) { - offset_index_writer->Abort(); - } - }); - PAIMON_RETURN_NOT_OK(FlushSegment(segment, offset_index_writer.get())); + PAIMON_RETURN_NOT_OK(FlushSegment(segment)); PAIMON_ASSIGN_OR_RAISE(CommitIncrement segment_increment, file_writer_->PrepareCommit(wait_compaction)); const DataIncrement& increment = segment_increment.GetNewFilesIncrement(); - // TODO(xinyu.lxy): Make offset-index writing follow data-file rolling boundaries so one - // deduplicate segment can produce multiple data files with one `.offset` sidecar per file. - if (increment.NewFiles().size() != 1 || !increment.DeletedFiles().empty() || + if (increment.NewFiles().empty() || !increment.DeletedFiles().empty() || !increment.ChangelogFiles().empty() || !segment_increment.GetCompactIncrement().IsEmpty() || segment_increment.GetCompactDeletionFile()) { - return Status::Invalid( - "one non-empty deduplicate segment must produce one append data file"); + return Status::Invalid("deduplicate segment must produce append data files only"); } - std::shared_ptr<DataFileMeta> data_file = increment.NewFiles().front(); - if (data_file->row_count != segment->GetRowCount()) { + int64_t remaining_rows = segment->GetRowCount(); + for (const std::shared_ptr<DataFileMeta>& data_file : increment.NewFiles()) { + if (!data_file || data_file->row_count <= 0 || data_file->row_count > remaining_rows) { + return Status::Invalid("deduplicate segment and data-file row counts differ"); + } + remaining_rows -= data_file->row_count; + new_files.push_back(data_file); + } + if (remaining_rows != 0) { return Status::Invalid("deduplicate segment and data-file row counts differ"); } - PAIMON_ASSIGN_OR_RAISE(std::string offset_index_file, - offset_index_writer->Finish(data_file)); - std::vector<std::optional<std::string>> extra_files = data_file->extra_files; - extra_files.emplace_back(std::move(offset_index_file)); - data_file = data_file->CopyWithExtraFiles(extra_files); - new_files.push_back(data_file); AppendAll(increment.NewIndexFiles(), &new_index_files); AppendAll(increment.DeletedIndexFiles(), &deleted_index_files); - offset_index_guard.Release(); } DataIncrement data_increment(std::move(new_files), {}, {}, std::move(new_index_files), @@ -269,11 +249,7 @@ } Status RealtimeDeduplicateWriter::FlushSegment( - const std::shared_ptr<RealtimeSegmentHandle>& segment, - DataFileKeyOffsetIndexWriter* offset_index_writer) { - if (!offset_index_writer) { - return Status::Invalid("deduplicate offset-index writer is null"); - } + const std::shared_ptr<RealtimeSegmentHandle>& segment) { PAIMON_ASSIGN_OR_RAISE(std::vector<std::unique_ptr<BatchReader>> readers, realtime_store_->CreateCommitReaders(segment)); ConcatBatchReader reader(std::move(readers), arrow_pool_); @@ -292,17 +268,6 @@ } std::shared_ptr<arrow::StructArray> data = checked_pointer_cast<arrow::StructArray>(imported); - std::shared_ptr<arrow::Array> raw_offsets = - data->GetFieldByName(SpecialFields::RealtimeOffset().Name()); - if (!raw_offsets || raw_offsets->type_id() != arrow::Type::INT64) { - return Status::Invalid("deduplicate commit batch has no realtime offset"); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr<arrow::StructArray> keys, - arrow::StructArray::Make({data->GetFieldByName(deduplicate_key_field_->name())}, - {deduplicate_key_field_})); - PAIMON_RETURN_NOT_OK(offset_index_writer->AddBatch( - keys, checked_pointer_cast<arrow::Int64Array>(raw_offsets))); emitted_rows += data->length(); auto output = std::make_unique<ArrowArray>();
diff --git a/src/paimon/core/realtime/realtime_deduplicate_writer.h b/src/paimon/core/realtime/realtime_deduplicate_writer.h index 2644a32..90d2551 100644 --- a/src/paimon/core/realtime/realtime_deduplicate_writer.h +++ b/src/paimon/core/realtime/realtime_deduplicate_writer.h
@@ -36,9 +36,7 @@ class AppendOnlyWriter; class ArrowDeduplicateRealtimeStore; class CoreOptions; -class DataFileKeyOffsetIndexWriter; class DataFilePathFactory; -class FileSystem; class MemoryPool; class RealtimeContext; class RealtimeSchemaLayout; @@ -77,24 +75,16 @@ std::shared_ptr<AppendOnlyWriter> file_writer, std::shared_ptr<RealtimeSchemaLayout> schema_layout, std::shared_ptr<arrow::Field> deduplicate_key_field, - std::shared_ptr<DataFilePathFactory> data_file_path_factory, - std::shared_ptr<FileSystem> file_system, int64_t next_offset, - std::shared_ptr<MemoryPool> memory_pool, - std::map<std::string, std::string> options); + int64_t next_offset, std::shared_ptr<MemoryPool> memory_pool); Status SealCurrentSegment(); - Status FlushSegment(const std::shared_ptr<RealtimeSegmentHandle>& segment, - DataFileKeyOffsetIndexWriter* offset_index_writer); + Status FlushSegment(const std::shared_ptr<RealtimeSegmentHandle>& segment); std::shared_ptr<arrow::MemoryPool> arrow_pool_; std::shared_ptr<ArrowDeduplicateRealtimeStore> realtime_store_; std::shared_ptr<AppendOnlyWriter> file_writer_; std::shared_ptr<RealtimeSchemaLayout> schema_layout_; std::shared_ptr<arrow::Field> deduplicate_key_field_; - std::shared_ptr<DataFilePathFactory> data_file_path_factory_; - std::shared_ptr<FileSystem> file_system_; - std::shared_ptr<MemoryPool> memory_pool_; - std::map<std::string, std::string> options_; std::vector<std::shared_ptr<RealtimeSegmentHandle>> sealed_segments_; int64_t next_offset_; bool has_building_data_ = false;
diff --git a/src/paimon/core/realtime/realtime_key_offset_index.cpp b/src/paimon/core/realtime/realtime_key_offset_index.cpp index 37ee94c..c57a41f 100644 --- a/src/paimon/core/realtime/realtime_key_offset_index.cpp +++ b/src/paimon/core/realtime/realtime_key_offset_index.cpp
@@ -19,63 +19,22 @@ #include "paimon/core/realtime/realtime_key_offset_index.h" +#include <algorithm> #include <optional> #include <utility> #include "arrow/api.h" #include "paimon/common/global_index/key_serializer.h" -#include "paimon/common/io/data_output_stream.h" -#include "paimon/common/predicate/literal_converter.h" +#include "paimon/common/io/offset_input_stream.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/realtime/key_offset_index_utils.h" #include "paimon/fs/file_system.h" #include "paimon/io/data_input_stream.h" #include "paimon/macros.h" -#include "paimon/memory/bytes.h" namespace paimon { -namespace { - -constexpr int64_t kOffsetIndexMagic = 0x5041494D4F464653LL; -constexpr int32_t kOffsetIndexVersion = 1; - -Result<std::vector<std::string>> EncodeKeyBatch( - const std::shared_ptr<arrow::StructArray>& keys, const std::shared_ptr<arrow::Field>& key_field, - const std::shared_ptr<KeySerializer>& key_serializer) { - if (!keys || keys->num_fields() != 1 || !keys->field(0)->type()->Equals(key_field->type())) { - return Status::Invalid( - "deduplicate key lookup must contain exactly the configured key field"); - } - if (keys->null_count() != 0 || keys->field(0)->null_count() != 0) { - return Status::Invalid("deduplicate key must not contain null"); - } - PAIMON_ASSIGN_OR_RAISE(std::vector<Literal> literals, - LiteralConverter::ConvertLiteralsFromArray(*keys->field(0), - /*own_data=*/true)); - std::vector<std::string> encoded; - encoded.reserve(literals.size()); - for (const Literal& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Bytes> bytes, key_serializer->Serialize(literal)); - std::string key; - if (bytes->size() > 0) { - key.append(bytes->data(), bytes->size()); - } - encoded.push_back(std::move(key)); - } - return encoded; -} - -Status WriteRaw(const std::shared_ptr<OutputStream>& output, const char* data, int64_t size) { - PAIMON_ASSIGN_OR_RAISE(int64_t written, output->Write(data, size)); - if (written != size) { - return Status::IOError("short write for realtime .offset index"); - } - return Status::OK(); -} - -} // namespace - Result<std::shared_ptr<MutableKeyOffsetIndex>> MutableKeyOffsetIndex::Create( const std::shared_ptr<arrow::Field>& key_field, const std::shared_ptr<MemoryPool>& memory_pool) { @@ -84,20 +43,42 @@ } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeySerializer> serializer, KeySerializer::Create(key_field->type(), memory_pool)); - return std::shared_ptr<MutableKeyOffsetIndex>( - new MutableKeyOffsetIndex(key_field, std::move(serializer), std::make_shared<OffsetMap>())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<OffsetMap> offsets, + OffsetMap::Create(/*initial_capacity=*/1024, memory_pool)); + return std::shared_ptr<MutableKeyOffsetIndex>(new MutableKeyOffsetIndex( + key_field, std::move(serializer), std::move(offsets), memory_pool)); } MutableKeyOffsetIndex::MutableKeyOffsetIndex(std::shared_ptr<arrow::Field> key_field, std::shared_ptr<KeySerializer> key_serializer, - std::shared_ptr<OffsetMap> offsets) + std::shared_ptr<OffsetMap> offsets, + std::shared_ptr<MemoryPool> memory_pool) : key_field_(std::move(key_field)), key_serializer_(std::move(key_serializer)), - offsets_(std::move(offsets)) {} + offsets_(std::move(offsets)), + memory_pool_(std::move(memory_pool)) {} -Result<std::vector<std::string>> MutableKeyOffsetIndex::EncodeKeys( +Result<std::vector<uint64_t>> MutableKeyOffsetIndex::HashKeys( const std::shared_ptr<arrow::StructArray>& keys) const { - return EncodeKeyBatch(keys, key_field_, key_serializer_); + return KeyOffsetIndexUtils::HashKeys(keys, key_field_, key_serializer_); +} + +Status MutableKeyOffsetIndex::EnsureMutable() { + if (offsets_.unique()) { + return Status::OK(); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<OffsetMap> copy, + OffsetMap::Create(std::max<uint64_t>(1, offsets_->Size()), memory_pool_)); + Status status = Status::OK(); + offsets_->ForEach([&](uint64_t hash, int64_t offset) { + if (status.ok() && offset >= 0) { + status = copy->FindAndInsert(hash, offset); + } + }); + PAIMON_RETURN_NOT_OK(status); + offsets_ = std::move(copy); + return Status::OK(); } Status MutableKeyOffsetIndex::Add(const std::shared_ptr<arrow::StructArray>& keys, @@ -105,40 +86,38 @@ if (!offsets || offsets->null_count() != 0 || !keys || keys->length() != offsets->length()) { return Status::Invalid("key-offset index input columns are not aligned"); } - PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> encoded, EncodeKeys(keys)); - // Copy on write to keep previously pinned snapshots immutable. - if (!offsets_.unique()) { - offsets_ = std::make_shared<OffsetMap>(*offsets_); - } + PAIMON_ASSIGN_OR_RAISE(std::vector<uint64_t> hashes, HashKeys(keys)); for (int64_t row = 0; row < offsets->length(); ++row) { - const int64_t offset = offsets->Value(row); - if (offset < 0) { + if (offsets->Value(row) < 0) { return Status::Invalid("realtime offset must not be negative"); } - (*offsets_)[encoded[row]] = offset; + } + PAIMON_RETURN_NOT_OK(EnsureMutable()); + for (int64_t row = 0; row < offsets->length(); ++row) { + PAIMON_RETURN_NOT_OK(offsets_->FindAndInsert(hashes[row], offsets->Value(row))); } return Status::OK(); } Status MutableKeyOffsetIndex::Erase(const std::shared_ptr<arrow::StructArray>& keys) { - PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> encoded, EncodeKeys(keys)); - if (!offsets_.unique()) { - offsets_ = std::make_shared<OffsetMap>(*offsets_); - } - for (const std::string& key : encoded) { - offsets_->erase(key); + PAIMON_ASSIGN_OR_RAISE(std::vector<uint64_t> hashes, HashKeys(keys)); + PAIMON_RETURN_NOT_OK(EnsureMutable()); + for (uint64_t hash : hashes) { + if (offsets_->Find(hash)) { + PAIMON_RETURN_NOT_OK(offsets_->FindAndInsert(hash, /*value=*/-1)); + } } return Status::OK(); } Result<RoaringBitmap64> MutableKeyOffsetIndex::LookupOffsets( const std::shared_ptr<arrow::StructArray>& keys) const { - PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> encoded, EncodeKeys(keys)); + PAIMON_ASSIGN_OR_RAISE(std::vector<uint64_t> hashes, HashKeys(keys)); RoaringBitmap64 result; - for (const std::string& key : encoded) { - auto iter = offsets_->find(key); - if (iter != offsets_->end()) { - result.Add(iter->second); + for (uint64_t hash : hashes) { + const std::optional<int64_t> offset = offsets_->Find(hash); + if (offset && offset.value() >= 0) { + result.Add(offset.value()); } } return result; @@ -146,7 +125,7 @@ std::shared_ptr<const KeyOffsetLookup> MutableKeyOffsetIndex::Snapshot() const { return std::shared_ptr<const KeyOffsetLookup>( - new MutableKeyOffsetIndex(key_field_, key_serializer_, offsets_)); + new MutableKeyOffsetIndex(key_field_, key_serializer_, offsets_, memory_pool_)); } Result<RoaringBitmap64> CompositeKeyOffsetLookup::LookupOffsets( @@ -162,93 +141,19 @@ return result; } -Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> DataFileKeyOffsetIndexWriter::Create( - const std::shared_ptr<arrow::Field>& key_field, - const std::shared_ptr<DataFilePathFactory>& path_factory, - const std::shared_ptr<FileSystem>& file_system, const std::shared_ptr<MemoryPool>& memory_pool, - const std::map<std::string, std::string>& options) { - (void)options; - if (!key_field || !path_factory || !file_system || !memory_pool) { - return Status::Invalid("key-offset writer is missing a required dependency"); - } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeySerializer> serializer, - KeySerializer::Create(key_field->type(), memory_pool)); - return std::unique_ptr<DataFileKeyOffsetIndexWriter>(new DataFileKeyOffsetIndexWriter( - key_field, path_factory, file_system, std::move(serializer))); -} - -DataFileKeyOffsetIndexWriter::DataFileKeyOffsetIndexWriter( - std::shared_ptr<arrow::Field> key_field, std::shared_ptr<DataFilePathFactory> path_factory, - std::shared_ptr<FileSystem> file_system, std::shared_ptr<KeySerializer> key_serializer) - : key_field_(std::move(key_field)), - path_factory_(std::move(path_factory)), - file_system_(std::move(file_system)), - key_serializer_(std::move(key_serializer)) {} - -Status DataFileKeyOffsetIndexWriter::AddBatch(const std::shared_ptr<arrow::StructArray>& keys, - const std::shared_ptr<arrow::Int64Array>& offsets) { - if (!keys || !offsets || keys->length() != offsets->length() || offsets->null_count() != 0) { - return Status::Invalid("key-offset writer input columns are not aligned"); - } - PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> encoded, - EncodeKeyBatch(keys, key_field_, key_serializer_)); - for (int64_t row = 0; row < offsets->length(); ++row) { - const int64_t offset = offsets->Value(row); - if (offset < 0) { - return Status::Invalid("realtime offset must not be negative"); - } - if (!offsets_.emplace(encoded[row], offset).second) { - return Status::Invalid( - "deduplicate data file contains more than one offset for a deduplicate key"); - } - } - return Status::OK(); -} - -Result<std::string> DataFileKeyOffsetIndexWriter::Finish( - const std::shared_ptr<DataFileMeta>& data_file) { - if (!data_file || offsets_.empty()) { - return Status::Invalid("cannot finish an empty key-offset index"); - } - const std::string file_name = data_file->file_name + DataFileKeyOffsetIndexReader::kFileSuffix; - output_path_ = path_factory_->ToAlignedPath(file_name, data_file); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<OutputStream> output, - file_system_->Create(output_path_, /*overwrite=*/false)); - DataOutputStream data_output(output); - PAIMON_RETURN_NOT_OK(data_output.WriteValue<int64_t>(kOffsetIndexMagic)); - PAIMON_RETURN_NOT_OK(data_output.WriteValue<int32_t>(kOffsetIndexVersion)); - PAIMON_RETURN_NOT_OK(data_output.WriteValue<int64_t>(offsets_.size())); - for (const auto& [key, offset] : offsets_) { - PAIMON_RETURN_NOT_OK(data_output.WriteValue<int32_t>(static_cast<int32_t>(key.size()))); - if (!key.empty()) { - PAIMON_RETURN_NOT_OK(WriteRaw(output, key.data(), key.size())); - } - PAIMON_RETURN_NOT_OK(data_output.WriteValue<int64_t>(offset)); - } - PAIMON_RETURN_NOT_OK(output->Flush()); - PAIMON_RETURN_NOT_OK(output->Close()); - return file_name; -} - -void DataFileKeyOffsetIndexWriter::Abort() { - if (!output_path_.empty()) { - [[maybe_unused]] Status status = file_system_->Delete(output_path_); - output_path_.clear(); - } -} - Result<std::shared_ptr<DataFileKeyOffsetIndexReader>> DataFileKeyOffsetIndexReader::Create( + uint32_t page_size, const std::shared_ptr<CacheManager>& cache_manager, const std::shared_ptr<arrow::Field>& key_field, const std::shared_ptr<DataFileMeta>& data_file, const std::shared_ptr<DataFilePathFactory>& path_factory, - const std::shared_ptr<FileSystem>& file_system, const std::shared_ptr<MemoryPool>& memory_pool, - const std::map<std::string, std::string>& options) { - (void)options; - if (!key_field || !data_file || !path_factory || !file_system || !memory_pool) { + const std::shared_ptr<FileSystem>& file_system, + const std::shared_ptr<MemoryPool>& memory_pool) { + if (!key_field || !data_file || !path_factory || !file_system || !memory_pool || + !cache_manager) { return Status::Invalid("key-offset reader is missing a required dependency"); } std::optional<std::string> file_name; for (const std::optional<std::string>& extra : data_file->extra_files) { - if (extra && StringUtils::EndsWith(extra.value(), kFileSuffix)) { + if (extra && StringUtils::EndsWith(extra.value(), KeyOffsetIndexUtils::kFileSuffix)) { if (file_name) { return Status::Invalid("data file contains multiple .offset sidecars"); } @@ -263,39 +168,23 @@ const std::string file_path = path_factory->ToAlignedPath(file_name.value(), data_file); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input, file_system->Open(file_path)); DataInputStream data_input(input); - PAIMON_ASSIGN_OR_RAISE(int64_t magic, data_input.ReadValue<int64_t>()); PAIMON_ASSIGN_OR_RAISE(int32_t version, data_input.ReadValue<int32_t>()); - if (magic != kOffsetIndexMagic || version != kOffsetIndexVersion) { - return Status::Invalid("unsupported realtime .offset index format: ", file_name.value()); + if (version != KeyOffsetIndexUtils::kCurrentVersion) { + return Status::Invalid("unsupported realtime .offset index version: ", version); } - PAIMON_ASSIGN_OR_RAISE(int64_t entry_count, data_input.ReadValue<int64_t>()); - if (entry_count < 0) { - return Status::Invalid("negative entry count in realtime .offset index"); + PAIMON_ASSIGN_OR_RAISE(int64_t length, input->Length()); + constexpr int64_t kVersionSize = sizeof(int32_t); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<OffsetInputStream> hash_table_input, + OffsetInputStream::Create(input, length - kVersionSize, kVersionSize, length)); + using Reader = DenseHashTableFileReader<uint64_t, int64_t>; + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<Reader> offsets, + Reader::Create(page_size, cache_manager, file_path, hash_table_input, memory_pool)); + if (data_file->row_count < 0 || + offsets->Size() != static_cast<uint64_t>(data_file->row_count)) { + return Status::Invalid("realtime .offset index entry count does not match data file"); } - auto offsets = std::make_shared<std::map<std::string, int64_t>>(); - for (int64_t entry = 0; entry < entry_count; ++entry) { - PAIMON_ASSIGN_OR_RAISE(int32_t key_size, data_input.ReadValue<int32_t>()); - if (key_size < 0) { - return Status::Invalid("invalid key size in realtime .offset index"); - } - std::string key(key_size, '\0'); - if (key_size > 0) { - PAIMON_RETURN_NOT_OK(data_input.Read(key.data(), key_size)); - } - PAIMON_ASSIGN_OR_RAISE(int64_t offset, data_input.ReadValue<int64_t>()); - if (offset < 0) { - return Status::Invalid("negative offset in realtime .offset index"); - } - if (!offsets->emplace(std::move(key), offset).second) { - return Status::Invalid("duplicate key in realtime .offset index"); - } - } - PAIMON_ASSIGN_OR_RAISE(int64_t position, data_input.GetPos()); - PAIMON_ASSIGN_OR_RAISE(int64_t length, data_input.Length()); - if (position != length) { - return Status::Invalid("trailing bytes in realtime .offset index"); - } - PAIMON_RETURN_NOT_OK(input->Close()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<KeySerializer> serializer, KeySerializer::Create(key_field->type(), memory_pool)); return std::shared_ptr<DataFileKeyOffsetIndexReader>( @@ -304,13 +193,16 @@ Result<RoaringBitmap64> DataFileKeyOffsetIndexReader::LookupOffsets( const std::shared_ptr<arrow::StructArray>& keys) const { - PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> encoded, - EncodeKeyBatch(keys, key_field_, key_serializer_)); + PAIMON_ASSIGN_OR_RAISE(std::vector<uint64_t> hashes, + KeyOffsetIndexUtils::HashKeys(keys, key_field_, key_serializer_)); RoaringBitmap64 result; - for (const std::string& key : encoded) { - auto iter = offsets_->find(key); - if (iter != offsets_->end()) { - result.Add(iter->second); + for (uint64_t hash : hashes) { + PAIMON_ASSIGN_OR_RAISE(std::optional<int64_t> offset, offsets_->Find(hash)); + if (offset) { + if (offset.value() < 0) { + return Status::Invalid("negative offset in realtime .offset index"); + } + result.Add(offset.value()); } } return result;
diff --git a/src/paimon/core/realtime/realtime_key_offset_index.h b/src/paimon/core/realtime/realtime_key_offset_index.h index 0c8e9f8..fe799d8 100644 --- a/src/paimon/core/realtime/realtime_key_offset_index.h +++ b/src/paimon/core/realtime/realtime_key_offset_index.h
@@ -19,13 +19,15 @@ #pragma once -#include <map> +#include <cstdint> #include <memory> #include <string> #include <utility> #include <vector> #include "arrow/type_fwd.h" +#include "paimon/common/utils/dense_hash_table_file_reader.h" +#include "paimon/common/utils/hash_map.h" #include "paimon/result.h" #include "paimon/utils/roaring_bitmap64.h" @@ -47,7 +49,8 @@ }; /// Copy-on-write in-memory implementation used by one building segment. Each key retains only its -/// latest live offset; overwrites replace the value and deletes erase the key. +/// latest live offset. Only 64-bit hashes are retained; hash collisions are not distinguished. +/// Deletes mark existing hashes with a negative offset until the segment is released. class MutableKeyOffsetIndex final : public KeyOffsetLookup { public: static Result<std::shared_ptr<MutableKeyOffsetIndex>> Create( @@ -65,18 +68,21 @@ std::shared_ptr<const KeyOffsetLookup> Snapshot() const; private: - using OffsetMap = std::map<std::string, int64_t>; + using OffsetMap = HashMap<uint64_t, int64_t>; MutableKeyOffsetIndex(std::shared_ptr<arrow::Field> key_field, std::shared_ptr<KeySerializer> key_serializer, - std::shared_ptr<OffsetMap> offsets); + std::shared_ptr<OffsetMap> offsets, + std::shared_ptr<MemoryPool> memory_pool); - Result<std::vector<std::string>> EncodeKeys( - const std::shared_ptr<arrow::StructArray>& keys) const; + Result<std::vector<uint64_t>> HashKeys(const std::shared_ptr<arrow::StructArray>& keys) const; + + Status EnsureMutable(); std::shared_ptr<arrow::Field> key_field_; std::shared_ptr<KeySerializer> key_serializer_; std::shared_ptr<OffsetMap> offsets_; + std::shared_ptr<MemoryPool> memory_pool_; }; /// Unions a fixed set of immutable lookups. The set itself is pinned by a writer lookup view. @@ -92,67 +98,31 @@ std::vector<std::shared_ptr<const KeyOffsetLookup>> lookups_; }; -/// Buffers one data file's surviving key/offset pairs and writes a temporary `.offset` sidecar. -/// Segment-level deduplication guarantees at most one offset per key in the resulting data file. -/// TODO(xinyu.lxy): This writer and its `.offset` sidecar format are experimental. The format will -/// be revised before production use. -class DataFileKeyOffsetIndexWriter { - public: - static Result<std::unique_ptr<DataFileKeyOffsetIndexWriter>> Create( - const std::shared_ptr<arrow::Field>& key_field, - const std::shared_ptr<DataFilePathFactory>& path_factory, - const std::shared_ptr<FileSystem>& file_system, - const std::shared_ptr<MemoryPool>& memory_pool, - const std::map<std::string, std::string>& options); - - Status AddBatch(const std::shared_ptr<arrow::StructArray>& keys, - const std::shared_ptr<arrow::Int64Array>& offsets); - - Result<std::string> Finish(const std::shared_ptr<DataFileMeta>& data_file); - - void Abort(); - - private: - DataFileKeyOffsetIndexWriter(std::shared_ptr<arrow::Field> key_field, - std::shared_ptr<DataFilePathFactory> path_factory, - std::shared_ptr<FileSystem> file_system, - std::shared_ptr<KeySerializer> key_serializer); - - std::shared_ptr<arrow::Field> key_field_; - std::shared_ptr<DataFilePathFactory> path_factory_; - std::shared_ptr<FileSystem> file_system_; - std::shared_ptr<KeySerializer> key_serializer_; - std::map<std::string, int64_t> offsets_; - std::string output_path_; -}; - /// Reads one temporary `.offset` sidecar attached to a data file. class DataFileKeyOffsetIndexReader final : public KeyOffsetLookup { public: - static constexpr const char* kFileSuffix = ".offset"; - static Result<std::shared_ptr<DataFileKeyOffsetIndexReader>> Create( + uint32_t page_size, const std::shared_ptr<CacheManager>& cache_manager, const std::shared_ptr<arrow::Field>& key_field, const std::shared_ptr<DataFileMeta>& data_file, const std::shared_ptr<DataFilePathFactory>& path_factory, const std::shared_ptr<FileSystem>& file_system, - const std::shared_ptr<MemoryPool>& memory_pool, - const std::map<std::string, std::string>& options); + const std::shared_ptr<MemoryPool>& memory_pool); Result<RoaringBitmap64> LookupOffsets( const std::shared_ptr<arrow::StructArray>& keys) const override; private: - DataFileKeyOffsetIndexReader(std::shared_ptr<arrow::Field> key_field, - std::shared_ptr<KeySerializer> key_serializer, - std::shared_ptr<const std::map<std::string, int64_t>> offsets) + DataFileKeyOffsetIndexReader( + std::shared_ptr<arrow::Field> key_field, std::shared_ptr<KeySerializer> key_serializer, + std::unique_ptr<DenseHashTableFileReader<uint64_t, int64_t>> offsets) : key_field_(std::move(key_field)), key_serializer_(std::move(key_serializer)), offsets_(std::move(offsets)) {} std::shared_ptr<arrow::Field> key_field_; std::shared_ptr<KeySerializer> key_serializer_; - std::shared_ptr<const std::map<std::string, int64_t>> offsets_; + std::unique_ptr<DenseHashTableFileReader<uint64_t, int64_t>> offsets_; }; } // namespace paimon
diff --git a/src/paimon/core/realtime/realtime_key_offset_index_test.cpp b/src/paimon/core/realtime/realtime_key_offset_index_test.cpp index c58c4f7..5275578 100644 --- a/src/paimon/core/realtime/realtime_key_offset_index_test.cpp +++ b/src/paimon/core/realtime/realtime_key_offset_index_test.cpp
@@ -27,12 +27,20 @@ #include "arrow/api.h" #include "arrow/ipc/json_simple.h" +#include "paimon/common/global_index/key_serializer.h" +#include "paimon/common/io/cache/cache_manager.h" +#include "paimon/common/io/cache/lru_cache.h" +#include "paimon/common/io/data_output_stream.h" +#include "paimon/common/table/special_fields.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/data_file_path_factory.h" +#include "paimon/core/realtime/data_file_key_offset_index_writer.h" +#include "paimon/core/realtime/key_offset_index_utils.h" #include "paimon/fs/file_system.h" #include "paimon/memory/memory_pool.h" #include "paimon/testing/utils/testharness.h" +#include "xxhash.h" // NOLINT(build/include_subdir) namespace paimon::test { namespace { @@ -48,6 +56,16 @@ arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), json).ValueOrDie()); } +std::shared_ptr<arrow::StructArray> MakeIndexBatch(const std::shared_ptr<arrow::Field>& key_field, + const std::string& json) { + return checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_( + {key_field, arrow::field(SpecialFields::RealtimeOffset().Name(), arrow::int64())}), + json) + .ValueOrDie()); +} + std::shared_ptr<DataFileMeta> MakeDataFile( const std::string& file_name, int64_t row_count, const std::vector<std::optional<std::string>>& extra_files = {}) { @@ -87,6 +105,14 @@ ASSERT_OK_AND_ASSIGN(RoaringBitmap64 pinned, snapshot->LookupOffsets(MakeKeys(key_field, R"([[1], [2], [3]])"))); AssertOffsets(pinned, {10, 20}); + + const std::shared_ptr<const KeyOffsetLookup> deleted_snapshot = index->Snapshot(); + ASSERT_OK(index->Add(MakeKeys(key_field, R"([[2]])"), MakeOffsets("[40]"))); + ASSERT_OK(index->Erase(MakeKeys(key_field, R"([[99]])"))); + ASSERT_OK_AND_ASSIGN(current, index->LookupOffsets(MakeKeys(key_field, R"([[1], [2]])"))); + AssertOffsets(current, {30, 40}); + ASSERT_OK_AND_ASSIGN(pinned, deleted_snapshot->LookupOffsets(MakeKeys(key_field, R"([[2]])"))); + ASSERT_TRUE(pinned.IsEmpty()); } TEST(RealtimeKeyOffsetIndexTest, TestMutableIndexValidatesInput) { @@ -127,6 +153,7 @@ } TEST(RealtimeKeyOffsetIndexTest, TestSidecarRoundTripAndAbort) { + auto cache_manager = std::make_shared<CacheManager>(128, /*high_priority_pool_ratio=*/0.0); std::unique_ptr<UniqueTestDirectory> directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); const std::shared_ptr<arrow::Field> key_field = arrow::field("id", arrow::int64()); @@ -134,25 +161,37 @@ ASSERT_OK(path_factory->Init(directory->Str(), /*format_identifier=*/"parquet", /*data_file_prefix=*/"data-", /*external_path_provider=*/nullptr)); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, - DataFileKeyOffsetIndexWriter::Create(key_field, path_factory, directory->GetFileSystem(), - GetDefaultPool(), /*options=*/{})); - ASSERT_OK(writer->AddBatch(MakeKeys(key_field, R"([[1], [2]])"), MakeOffsets("[7, 9]"))); - ASSERT_NOK_WITH_MSG(writer->AddBatch(MakeKeys(key_field, R"([[2]])"), MakeOffsets("[11]")), + ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, + DataFileKeyOffsetIndexWriter::Create(key_field, directory->GetFileSystem(), + GetDefaultPool())); + ASSERT_OK(writer->AddBatch(MakeIndexBatch(key_field, R"([[1, 7], [2, 9]])"))); + ASSERT_NOK_WITH_MSG(writer->AddBatch(MakeIndexBatch(key_field, R"([[2, 11]])")), "more than one offset for a deduplicate key"); std::shared_ptr<DataFileMeta> data_file = MakeDataFile("data-0.parquet", /*row_count=*/2); - ASSERT_OK_AND_ASSIGN(std::string sidecar, writer->Finish(data_file)); + ASSERT_OK_AND_ASSIGN(std::string sidecar, + writer->Finish(path_factory->ToPath(data_file), data_file->row_count)); ASSERT_EQ("data-0.parquet.offset", sidecar); data_file->extra_files.emplace_back(sidecar); - ASSERT_OK_AND_ASSIGN( - std::shared_ptr<DataFileKeyOffsetIndexReader> reader, - DataFileKeyOffsetIndexReader::Create(key_field, data_file, path_factory, - directory->GetFileSystem(), GetDefaultPool(), {})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create( + /*page_size=*/32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool())); + const std::shared_ptr<LruCache> cache = + checked_pointer_cast<LruCache>(cache_manager->IndexCache()); + ASSERT_EQ(0, cache->GetCurrentWeight()); ASSERT_OK_AND_ASSIGN(RoaringBitmap64 result, reader->LookupOffsets(MakeKeys(key_field, R"([[2], [3], [1]])"))); AssertOffsets(result, {7, 9}); + ASSERT_GT(cache->GetCurrentWeight(), 0); + ASSERT_LE(cache->GetCurrentWeight(), 128); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataFileKeyOffsetIndexReader> reopened, + DataFileKeyOffsetIndexReader::Create( + /*page_size=*/32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(result, reopened->LookupOffsets(MakeKeys(key_field, R"([[1], [2]])"))); + AssertOffsets(result, {7, 9}); const std::string sidecar_path = path_factory->ToAlignedPath(sidecar, data_file); ASSERT_OK_AND_ASSIGN(bool exists, directory->GetFileSystem()->Exists(sidecar_path)); @@ -160,11 +199,10 @@ writer->Abort(); ASSERT_OK_AND_ASSIGN(exists, directory->GetFileSystem()->Exists(sidecar_path)); ASSERT_FALSE(exists); - ASSERT_OK_AND_ASSIGN(result, reader->LookupOffsets(MakeKeys(key_field, R"([[1], [2]])"))); - AssertOffsets(result, {7, 9}); } TEST(RealtimeKeyOffsetIndexTest, TestSidecarRoundTripWithEmptyStringKey) { + auto cache_manager = std::make_shared<CacheManager>(128, /*high_priority_pool_ratio=*/0.0); std::unique_ptr<UniqueTestDirectory> directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); const std::shared_ptr<arrow::Field> key_field = arrow::field("id", arrow::utf8()); @@ -172,19 +210,19 @@ ASSERT_OK(path_factory->Init(directory->Str(), /*format_identifier=*/"parquet", /*data_file_prefix=*/"data-", /*external_path_provider=*/nullptr)); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, - DataFileKeyOffsetIndexWriter::Create(key_field, path_factory, directory->GetFileSystem(), - GetDefaultPool(), /*options=*/{})); - ASSERT_OK(writer->AddBatch(MakeKeys(key_field, R"([[""], ["alpha"]])"), MakeOffsets("[7, 9]"))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, + DataFileKeyOffsetIndexWriter::Create(key_field, directory->GetFileSystem(), + GetDefaultPool())); + ASSERT_OK(writer->AddBatch(MakeIndexBatch(key_field, R"([["", 7], ["alpha", 9]])"))); std::shared_ptr<DataFileMeta> data_file = MakeDataFile("data-0.parquet", /*row_count=*/2); - ASSERT_OK_AND_ASSIGN(std::string sidecar, writer->Finish(data_file)); + ASSERT_OK_AND_ASSIGN(std::string sidecar, + writer->Finish(path_factory->ToPath(data_file), data_file->row_count)); data_file->extra_files.emplace_back(sidecar); - ASSERT_OK_AND_ASSIGN( - std::shared_ptr<DataFileKeyOffsetIndexReader> reader, - DataFileKeyOffsetIndexReader::Create(key_field, data_file, path_factory, - directory->GetFileSystem(), GetDefaultPool(), {})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create( + /*page_size=*/32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool())); ASSERT_OK_AND_ASSIGN(RoaringBitmap64 empty_key_offset, reader->LookupOffsets(MakeKeys(key_field, R"([[""]])"))); AssertOffsets(empty_key_offset, {7}); @@ -196,7 +234,47 @@ ASSERT_TRUE(missing_key_offset.IsEmpty()); } +TEST(RealtimeKeyOffsetIndexTest, TestSidecarAcceptsDictionaryEncodedLogicalBatch) { + auto directory = UniqueTestDirectory::Create(); + ASSERT_TRUE(directory); + const std::shared_ptr<arrow::Field> key_field = arrow::field("id", arrow::utf8()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, + DataFileKeyOffsetIndexWriter::Create(key_field, directory->GetFileSystem(), + GetDefaultPool())); + const std::shared_ptr<arrow::Array> indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[1, 0]").ValueOrDie(); + const std::shared_ptr<arrow::Array> dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b"])").ValueOrDie(); + const std::shared_ptr<arrow::Array> encoded = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::utf8()), + indices, dictionary) + .ValueOrDie(); + const std::shared_ptr<arrow::StructArray> logical_batch = + arrow::StructArray::Make( + {encoded, MakeOffsets("[100, 200]")}, + std::vector<std::string>{"id", SpecialFields::RealtimeOffset().Name()}) + .ValueOrDie(); + ASSERT_OK(writer->AddBatch(logical_batch)); + auto path_factory = std::make_shared<DataFilePathFactory>(); + ASSERT_OK(path_factory->Init(directory->Str(), "parquet", "data-", nullptr)); + std::shared_ptr<DataFileMeta> data_file = MakeDataFile("dictionary.parquet", 2); + ASSERT_OK_AND_ASSIGN(std::string sidecar, + writer->Finish(path_factory->ToPath(data_file), data_file->row_count)); + data_file = data_file->CopyWithExtraFiles({sidecar}); + auto cache_manager = std::make_shared<CacheManager>(128, 0.0); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create(32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(RoaringBitmap64 offsets, + reader->LookupOffsets(MakeKeys(key_field, R"([["a"]])"))); + AssertOffsets(offsets, {200}); + ASSERT_OK_AND_ASSIGN(offsets, reader->LookupOffsets(MakeKeys(key_field, R"([["b"]])"))); + AssertOffsets(offsets, {100}); +} + TEST(RealtimeKeyOffsetIndexTest, TestSidecarValidation) { + auto cache_manager = std::make_shared<CacheManager>(128, /*high_priority_pool_ratio=*/0.0); std::unique_ptr<UniqueTestDirectory> directory = UniqueTestDirectory::Create(); ASSERT_NE(nullptr, directory); const std::shared_ptr<arrow::Field> key_field = arrow::field("id", arrow::int64()); @@ -205,26 +283,71 @@ /*data_file_prefix=*/"data-", /*external_path_provider=*/nullptr)); - ASSERT_NOK_WITH_MSG(DataFileKeyOffsetIndexWriter::Create( - key_field, nullptr, directory->GetFileSystem(), GetDefaultPool(), {}), + ASSERT_NOK_WITH_MSG(DataFileKeyOffsetIndexWriter::Create(key_field, nullptr, GetDefaultPool()), "missing a required dependency"); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, - DataFileKeyOffsetIndexWriter::Create(key_field, path_factory, directory->GetFileSystem(), - GetDefaultPool(), /*options=*/{})); - ASSERT_NOK_WITH_MSG(writer->Finish(MakeDataFile("empty.parquet", /*row_count=*/0)), + ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataFileKeyOffsetIndexWriter> writer, + DataFileKeyOffsetIndexWriter::Create(key_field, directory->GetFileSystem(), + GetDefaultPool())); + ASSERT_NOK_WITH_MSG(writer->Finish(path_factory->ToPath("empty.parquet"), /*row_count=*/0), "cannot finish an empty key-offset index"); ASSERT_NOK_WITH_MSG(DataFileKeyOffsetIndexReader::Create( - key_field, MakeDataFile("missing.parquet", /*row_count=*/1), - path_factory, directory->GetFileSystem(), GetDefaultPool(), {}), + /*page_size=*/32, cache_manager, key_field, + MakeDataFile("missing.parquet", /*row_count=*/1), path_factory, + directory->GetFileSystem(), GetDefaultPool()), "does not contain the required .offset sidecar"); ASSERT_NOK_WITH_MSG( DataFileKeyOffsetIndexReader::Create( - key_field, + /*page_size=*/32, cache_manager, key_field, MakeDataFile("duplicate.parquet", /*row_count=*/1, {"first.offset", "second.offset"}), - path_factory, directory->GetFileSystem(), GetDefaultPool(), {}), + path_factory, directory->GetFileSystem(), GetDefaultPool()), "contains multiple .offset sidecars"); + + ASSERT_OK(writer->AddBatch(MakeIndexBatch(key_field, R"([[1, 0]])"))); + std::shared_ptr<DataFileMeta> data_file = MakeDataFile("one.parquet", /*row_count=*/1); + ASSERT_OK_AND_ASSIGN(std::string sidecar, + writer->Finish(path_factory->ToPath(data_file), data_file->row_count)); + data_file = MakeDataFile("one.parquet", /*row_count=*/2, {sidecar}); + ASSERT_NOK_WITH_MSG( + DataFileKeyOffsetIndexReader::Create(32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool()), + "entry count does not match data file"); +} + +TEST(RealtimeKeyOffsetIndexTest, TestSidecarRejectsNegativeOffsetOnLookup) { + std::unique_ptr<UniqueTestDirectory> directory = UniqueTestDirectory::Create(); + ASSERT_NE(nullptr, directory); + const std::shared_ptr<arrow::Field> key_field = arrow::field("id", arrow::int64()); + auto path_factory = std::make_shared<DataFilePathFactory>(); + ASSERT_OK(path_factory->Init(directory->Str(), /*format_identifier=*/"parquet", + /*data_file_prefix=*/"data-", /*external_path_provider=*/nullptr)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<KeySerializer> serializer, + KeySerializer::Create(key_field->type(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Bytes> serialized, + serializer->Serialize(Literal(int64_t{1}))); + using HashTable = DenseHashTable<uint64_t, int64_t>; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<HashTable> hash_table, + HashTable::Create(1, /*occupancy_pct=*/80, GetDefaultPool())); + ASSERT_OK(hash_table->Insert(XXH64(serialized->data(), serialized->size(), /*seed=*/0), -1)); + const std::shared_ptr<DataFileMeta> data_file = + MakeDataFile("negative.parquet", /*row_count=*/1, {"negative.parquet.offset"}); + const std::string path = path_factory->ToAlignedPath("negative.parquet.offset", data_file); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> output, + directory->GetFileSystem()->Create(path, /*overwrite=*/false)); + DataOutputStream data_output(output); + ASSERT_OK(data_output.WriteValue<int32_t>(KeyOffsetIndexUtils::kCurrentVersion)); + ASSERT_OK(hash_table->WriteTo(output.get())); + ASSERT_OK(output->Flush()); + ASSERT_OK(output->Close()); + + auto cache_manager = std::make_shared<CacheManager>(128, /*high_priority_pool_ratio=*/0.0); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create(32, cache_manager, key_field, data_file, path_factory, + directory->GetFileSystem(), GetDefaultPool())); + // Opening checks the header only. Invalid values are rejected when their buckets are read. + ASSERT_NOK_WITH_MSG(reader->LookupOffsets(MakeKeys(key_field, R"([[1]])")), + "negative offset in realtime .offset index"); } } // namespace paimon::test
diff --git a/src/paimon/core/realtime/realtime_offset_file_index_lookup.cpp b/src/paimon/core/realtime/realtime_offset_file_index_lookup.cpp index 83f39ac..d9dacd7 100644 --- a/src/paimon/core/realtime/realtime_offset_file_index_lookup.cpp +++ b/src/paimon/core/realtime/realtime_offset_file_index_lookup.cpp
@@ -24,6 +24,7 @@ #include "arrow/api.h" #include "fmt/format.h" +#include "paimon/common/io/cache/cache_manager.h" #include "paimon/common/table/special_fields.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" @@ -55,10 +56,15 @@ return Status::Invalid("invalid data file for realtime offset lookup"); } } - return std::shared_ptr<RealtimeOffsetFileIndexLookup>( - new RealtimeOffsetFileIndexLookup(data_schema, data_schema_id, deduplicate_key_field, - std::vector<std::shared_ptr<DataFileMeta>>(data_files), - path_factory, file_system, memory_pool, options.ToMap())); + if (options.GetCachePageSize() <= 0 || options.GetLookupCacheMaxMemory() < 0) { + return Status::Invalid("invalid realtime key-offset cache options"); + } + auto cache_manager = std::make_shared<CacheManager>(options.GetLookupCacheMaxMemory(), + options.GetLookupCacheHighPrioPoolRatio()); + return std::shared_ptr<RealtimeOffsetFileIndexLookup>(new RealtimeOffsetFileIndexLookup( + data_schema, data_schema_id, deduplicate_key_field, + std::vector<std::shared_ptr<DataFileMeta>>(data_files), path_factory, file_system, + options.ToMap(), options.GetCachePageSize(), std::move(cache_manager), memory_pool)); } RealtimeOffsetFileIndexLookup::RealtimeOffsetFileIndexLookup( @@ -66,7 +72,8 @@ std::shared_ptr<arrow::Field> deduplicate_key_field, std::vector<std::shared_ptr<DataFileMeta>> data_files, std::shared_ptr<DataFilePathFactory> path_factory, std::shared_ptr<FileSystem> file_system, - std::shared_ptr<MemoryPool> memory_pool, std::map<std::string, std::string> options) + std::map<std::string, std::string> options, uint32_t cache_page_size, + std::shared_ptr<CacheManager> cache_manager, std::shared_ptr<MemoryPool> memory_pool) : data_schema_(std::move(data_schema)), data_schema_id_(data_schema_id), deduplicate_key_field_(std::move(deduplicate_key_field)), @@ -74,7 +81,9 @@ path_factory_(std::move(path_factory)), file_system_(std::move(file_system)), memory_pool_(std::move(memory_pool)), - options_(std::move(options)) {} + options_(std::move(options)), + cache_page_size_(cache_page_size), + cache_manager_(std::move(cache_manager)) {} Result<std::map<std::string, RoaringBitmap32>> RealtimeOffsetFileIndexLookup::LookupFilePositions( const RoaringBitmap64& offsets) const { @@ -158,10 +167,10 @@ std::vector<std::shared_ptr<const KeyOffsetLookup>> readers; readers.reserve(data_files_.size()); for (const std::shared_ptr<DataFileMeta>& file : data_files_) { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr<DataFileKeyOffsetIndexReader> reader, - DataFileKeyOffsetIndexReader::Create(deduplicate_key_field_, file, path_factory_, - file_system_, memory_pool_, options_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataFileKeyOffsetIndexReader> reader, + DataFileKeyOffsetIndexReader::Create( + cache_page_size_, cache_manager_, deduplicate_key_field_, file, + path_factory_, file_system_, memory_pool_)); readers.push_back(std::move(reader)); } return std::shared_ptr<const KeyOffsetLookup>(new CompositeKeyOffsetLookup(std::move(readers))); @@ -174,10 +183,10 @@ return Status::Invalid("invalid data file for realtime offset lookup"); } } - return std::shared_ptr<RealtimeOffsetFileIndexLookup>( - new RealtimeOffsetFileIndexLookup(data_schema_, data_schema_id_, deduplicate_key_field_, - std::vector<std::shared_ptr<DataFileMeta>>(data_files), - path_factory_, file_system_, memory_pool_, options_)); + return std::shared_ptr<RealtimeOffsetFileIndexLookup>(new RealtimeOffsetFileIndexLookup( + data_schema_, data_schema_id_, deduplicate_key_field_, + std::vector<std::shared_ptr<DataFileMeta>>(data_files), path_factory_, file_system_, + options_, cache_page_size_, cache_manager_, memory_pool_)); } } // namespace paimon
diff --git a/src/paimon/core/realtime/realtime_offset_file_index_lookup.h b/src/paimon/core/realtime/realtime_offset_file_index_lookup.h index fa21437..8ff4531 100644 --- a/src/paimon/core/realtime/realtime_offset_file_index_lookup.h +++ b/src/paimon/core/realtime/realtime_offset_file_index_lookup.h
@@ -35,6 +35,7 @@ namespace paimon { +class CacheManager; class CoreOptions; class DataFilePathFactory; class FileSystem; @@ -67,14 +68,13 @@ } private: - RealtimeOffsetFileIndexLookup(std::shared_ptr<arrow::Schema> data_schema, - int64_t data_schema_id, - std::shared_ptr<arrow::Field> deduplicate_key_field, - std::vector<std::shared_ptr<DataFileMeta>> data_files, - std::shared_ptr<DataFilePathFactory> path_factory, - std::shared_ptr<FileSystem> file_system, - std::shared_ptr<MemoryPool> memory_pool, - std::map<std::string, std::string> options); + RealtimeOffsetFileIndexLookup( + std::shared_ptr<arrow::Schema> data_schema, int64_t data_schema_id, + std::shared_ptr<arrow::Field> deduplicate_key_field, + std::vector<std::shared_ptr<DataFileMeta>> data_files, + std::shared_ptr<DataFilePathFactory> path_factory, std::shared_ptr<FileSystem> file_system, + std::map<std::string, std::string> options, uint32_t cache_page_size, + std::shared_ptr<CacheManager> cache_manager, std::shared_ptr<MemoryPool> memory_pool); std::shared_ptr<arrow::Schema> data_schema_; int64_t data_schema_id_; @@ -84,6 +84,9 @@ std::shared_ptr<FileSystem> file_system_; std::shared_ptr<MemoryPool> memory_pool_; std::map<std::string, std::string> options_; + uint32_t cache_page_size_; + // One budget for all sidecars and all pinned committed snapshots of this bucket. + std::shared_ptr<CacheManager> cache_manager_; }; } // namespace paimon
diff --git a/test/inte/realtime_write_inte_test.cpp b/test/inte/realtime_write_inte_test.cpp index 074dc9b..5de6c9f 100644 --- a/test/inte/realtime_write_inte_test.cpp +++ b/test/inte/realtime_write_inte_test.cpp
@@ -1291,6 +1291,84 @@ FinalizeCommitAndCheck(writer.get(), /*realtime_commits=*/{}, /*prepare_identifier=*/0, rows); } +TEST_F(RealtimeWriteInteTest, TestRealtimeDeduplicateRollingRestoresEveryOffsetSidecar) { + options_[Options::TARGET_FILE_ROW_NUM] = "2"; + ConfigureRealtimeDeduplicateTable(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> writer, + CreateRealtimeWriter(realtime_context)); + for (int64_t id = 1; id <= 5; id += 2) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch, + MakeBatch({Row{id, "old", "p0"}, Row{id + 1, "old", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(batch))); + } + const std::vector<Row> first_expected = {{1, "old", "p0"}, {2, "old", "p0"}, {3, "old", "p0"}, + {4, "old", "p0"}, {5, "old", "p0"}, {6, "old", "p0"}}; + ASSERT_OK_AND_ASSIGN(std::vector<Row> rows_before_prepare, ReadRows(realtime_context)); + ASSERT_EQ(first_expected, rows_before_prepare); + ASSERT_OK_AND_ASSIGN(std::vector<RealtimeCommitProgress> progress, + writer->PrepareCommitWithProgress(/*commit_identifier=*/0)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(0, 6), progress[0].offset_range); + std::shared_ptr<CommitMessageImpl> message = + std::dynamic_pointer_cast<CommitMessageImpl>(progress[0].commit_message); + ASSERT_NE(nullptr, message); + const std::vector<std::shared_ptr<DataFileMeta>>& files = + message->GetNewFilesIncrement().NewFiles(); + ASSERT_EQ(3, files.size()); + for (const std::shared_ptr<DataFileMeta>& file : files) { + ASSERT_EQ(2, file->row_count); + ASSERT_EQ(1, std::count(file->extra_files.begin(), file->extra_files.end(), + std::optional<std::string>(file->file_name + ".offset"))); + } + ASSERT_OK_AND_ASSIGN(std::vector<Row> rows_after_prepare, ReadRows(realtime_context)); + ASSERT_EQ(first_expected, rows_after_prepare); + ASSERT_OK_AND_ASSIGN(int64_t snapshot_id, Commit(progress, /*commit_identifier=*/0)); + ASSERT_OK_AND_ASSIGN(std::vector<Row> rows_before_refresh, ReadRows(realtime_context)); + ASSERT_EQ(first_expected, rows_before_refresh); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(std::vector<Row> rows_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(first_expected, rows_after_refresh); + ASSERT_OK(writer->Close()); + writer.reset(); + realtime_context.reset(); + + // Restart lookup from disk and overwrite one key from each of the three rolled files. + ASSERT_OK_AND_ASSIGN(realtime_context, RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(writer, CreateRealtimeWriter(realtime_context)); + ASSERT_OK_AND_ASSIGN(std::vector<Row> restored_rows, ReadRows(realtime_context)); + ASSERT_EQ(first_expected, restored_rows); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> update, + MakeBatch({Row{1, "new", "p0"}, Row{3, "new", "p0"}, Row{5, "new", "p0"}}, + /*partitioned=*/false)); + ASSERT_OK(writer->Write(std::move(update))); + const std::vector<Row> expected = {{2, "old", "p0"}, {4, "old", "p0"}, {6, "old", "p0"}, + {1, "new", "p0"}, {3, "new", "p0"}, {5, "new", "p0"}}; + ASSERT_OK_AND_ASSIGN(rows_before_prepare, ReadRows(realtime_context)); + ASSERT_EQ(expected, rows_before_prepare); + ASSERT_OK_AND_ASSIGN(progress, writer->PrepareCommitWithProgress(/*commit_identifier=*/1)); + ASSERT_EQ(1, progress.size()); + ASSERT_EQ(OffsetRange(6, 9), progress[0].offset_range); + ASSERT_NE(nullptr, progress[0].deleted_offsets); + ASSERT_EQ(3, progress[0].deleted_offsets->Cardinality()); + for (int64_t offset : {0, 2, 4}) { + ASSERT_TRUE(progress[0].deleted_offsets->Contains(offset)); + } + ASSERT_OK_AND_ASSIGN(rows_after_prepare, ReadRows(realtime_context)); + ASSERT_EQ(expected, rows_after_prepare); + ASSERT_OK_AND_ASSIGN(snapshot_id, Commit(progress, /*commit_identifier=*/1)); + ASSERT_OK_AND_ASSIGN(rows_before_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected, rows_before_refresh); + ASSERT_OK(writer->RefreshCommittedSnapshot(snapshot_id)); + ASSERT_OK_AND_ASSIGN(rows_after_refresh, ReadRows(realtime_context)); + ASSERT_EQ(expected, rows_after_refresh); + ASSERT_OK(writer->Close()); + ASSERT_OK_AND_ASSIGN(std::vector<Row> rows, ReadRows()); + ASSERT_EQ(expected, rows); +} + TEST_F(RealtimeWriteInteTest, TestRealtimeDeduplicateWritesDeletionVector) { ConfigureRealtimeDeduplicateTable();