diff --git a/docs/source/user_guide/manifest_cache.rst b/docs/source/user_guide/manifest_cache.rst index f6dd034af..6b45fe00c 100644 --- a/docs/source/user_guide/manifest_cache.rst +++ b/docs/source/user_guide/manifest_cache.rst @@ -102,7 +102,8 @@ concurrent cache misses. Cache implementations should use ``CacheValue::GetMemoryUsage()`` for admission and eviction accounting. For manifest IPC, this reports the retained Arrow -buffer capacity, including unused space from growth. ``GetSegment().Size()`` -continues to describe the valid IPC byte length. The built-in ``LruCache`` uses +buffer capacity, including unused space from growth. +``CacheValue::Get(value)->Size()`` describes the valid IPC +byte length. The built-in ``LruCache`` uses the memory usage value; existing cache values without an explicit allocation size continue to charge their segment length. diff --git a/include/paimon/cache/cache.h b/include/paimon/cache/cache.h index e05b555e5..e52a2fbe9 100644 --- a/include/paimon/cache/cache.h +++ b/include/paimon/cache/cache.h @@ -18,11 +18,13 @@ #pragma once +#include #include #include #include #include #include +#include #include "paimon/memory/memory_segment.h" #include "paimon/result.h" @@ -36,6 +38,7 @@ enum class CacheKind { MANIFEST, DATA_FILE_FOOTER, SNAPSHOT_LIVE_MANIFEST, + INDEX_READER, }; class PAIMON_EXPORT CacheKey { @@ -44,6 +47,8 @@ class PAIMON_EXPORT CacheKey { int32_t length, bool is_index); static std::shared_ptr ForKind(const std::string& file_path, int64_t position, int32_t length, CacheKind kind); + static std::shared_ptr ForIndexReader(const std::string& uri, + const std::string& column_name); public: virtual ~CacheKey() = default; @@ -88,28 +93,84 @@ class PAIMON_EXPORT Cache { class PAIMON_EXPORT CacheValue { public: - /// Charge the segment's logical size when no separate allocation size is supplied. - CacheValue(const MemorySegment& segment, CacheCallback callback); + virtual ~CacheValue(); - /// @param memory_usage Retained allocation size in bytes, including unused capacity. - /// The charge is at least the segment's logical size. - CacheValue(const MemorySegment& segment, CacheCallback callback, int64_t memory_usage); + /// Returns the typed payload, retaining its owner; null for a null value or type mismatch. + template + static std::shared_ptr Get(const std::shared_ptr& value); - ~CacheValue(); - - const MemorySegment& GetSegment() const; - - /// Bytes to charge against the cache budget. May exceed GetSegment().Size(). + /// Bytes charged against the cache budget. int64_t GetMemoryUsage() const; void OnEvict(const std::shared_ptr& key) const; - bool operator==(const CacheValue& other) const; + virtual bool operator==(const CacheValue& other) const = 0; + + protected: + CacheValue(CacheCallback callback, int64_t memory_usage); private: - MemorySegment segment_; CacheCallback callback_; int64_t memory_usage_; }; +template +class PAIMON_EXPORT TypedCacheValue : public CacheValue { + protected: + using CacheValue::CacheValue; + + private: + friend class CacheValue; + virtual T* GetValuePointer() = 0; +}; + +template +std::shared_ptr CacheValue::Get(const std::shared_ptr& value) { + auto* typed_value = dynamic_cast*>(value.get()); + if (typed_value) { + if (T* payload = typed_value->GetValuePointer()) { + return std::shared_ptr(value, payload); + } + } + return nullptr; +} + +class PAIMON_EXPORT MemorySegmentCacheValue final : public TypedCacheValue { + public: + MemorySegmentCacheValue(const MemorySegment& segment, CacheCallback callback); + + /// Charges retained allocation size, at least the segment's logical size. + MemorySegmentCacheValue(const MemorySegment& segment, CacheCallback callback, + int64_t memory_usage); + + bool operator==(const CacheValue& other) const override; + + private: + const MemorySegment* GetValuePointer() override { + return &segment_; + } + + MemorySegment segment_; +}; + +template +class PAIMON_EXPORT ObjectCacheValue final : public TypedCacheValue { + public: + ObjectCacheValue(std::shared_ptr value, CacheCallback callback, int64_t memory_usage) + : TypedCacheValue(std::move(callback), std::max(1, memory_usage)), + value_(std::move(value)) {} + + bool operator==(const CacheValue& other) const override { + const auto* rhs = dynamic_cast*>(&other); + return rhs && value_ == rhs->value_; + } + + private: + T* GetValuePointer() override { + return value_.get(); + } + + std::shared_ptr value_; +}; + } // namespace paimon diff --git a/src/paimon/common/io/cache/cache.cpp b/src/paimon/common/io/cache/cache.cpp index c4c9e5418..dc8ecb6b5 100644 --- a/src/paimon/common/io/cache/cache.cpp +++ b/src/paimon/common/io/cache/cache.cpp @@ -24,20 +24,11 @@ namespace paimon { -CacheValue::CacheValue(const MemorySegment& segment, CacheCallback callback) - : CacheValue(segment, std::move(callback), segment.Size()) {} - -CacheValue::CacheValue(const MemorySegment& segment, CacheCallback callback, int64_t memory_usage) - : segment_(segment), - callback_(std::move(callback)), - memory_usage_(std::max(segment.Size(), memory_usage)) {} +CacheValue::CacheValue(CacheCallback callback, int64_t memory_usage) + : callback_(std::move(callback)), memory_usage_(memory_usage) {} CacheValue::~CacheValue() = default; -const MemorySegment& CacheValue::GetSegment() const { - return segment_; -} - int64_t CacheValue::GetMemoryUsage() const { return memory_usage_; } @@ -48,11 +39,19 @@ void CacheValue::OnEvict(const std::shared_ptr& key) const { } } -bool CacheValue::operator==(const CacheValue& other) const { - if (this == &other) { - return true; - } - return segment_ == other.segment_; +MemorySegmentCacheValue::MemorySegmentCacheValue(const MemorySegment& segment, + CacheCallback callback) + : MemorySegmentCacheValue(segment, std::move(callback), segment.Size()) {} + +MemorySegmentCacheValue::MemorySegmentCacheValue(const MemorySegment& segment, + CacheCallback callback, int64_t memory_usage) + : TypedCacheValue(std::move(callback), + std::max(segment.Size(), memory_usage)), + segment_(segment) {} + +bool MemorySegmentCacheValue::operator==(const CacheValue& other) const { + const auto* rhs = dynamic_cast(&other); + return rhs && segment_ == rhs->segment_; } } // namespace paimon diff --git a/src/paimon/common/io/cache/cache_key.cpp b/src/paimon/common/io/cache/cache_key.cpp index 998d387f4..1dbd1aa3c 100644 --- a/src/paimon/common/io/cache/cache_key.cpp +++ b/src/paimon/common/io/cache/cache_key.cpp @@ -18,6 +18,8 @@ #include "paimon/common/io/cache/cache_key.h" +#include "fmt/format.h" + namespace paimon { std::shared_ptr CacheKey::ForPosition(const std::string& file_path, int64_t position, @@ -33,6 +35,13 @@ std::shared_ptr CacheKey::ForKind(const std::string& file_path, int64_ return key; } +std::shared_ptr CacheKey::ForIndexReader(const std::string& uri, + const std::string& column_name) { + std::string identity = + fmt::format("index-reader:{}:{}:{}:{}", uri.size(), uri, column_name.size(), column_name); + return ForKind(identity, 0, 0, CacheKind::INDEX_READER); +} + std::shared_ptr SnapshotLiveManifestEntriesCacheKey::ForExplicit( const std::string& table_path, const std::string& branch, int32_t bucket) { return std::shared_ptr(new SnapshotLiveManifestEntriesCacheKey( diff --git a/src/paimon/common/io/cache/cache_manager.cpp b/src/paimon/common/io/cache/cache_manager.cpp index 3e5c049b5..b3192b308 100644 --- a/src/paimon/common/io/cache/cache_manager.cpp +++ b/src/paimon/common/io/cache/cache_manager.cpp @@ -28,10 +28,14 @@ Result CacheManager::GetPage( auto supplier = [&](const std::shared_ptr& key) -> Result> { PAIMON_ASSIGN_OR_RAISE(MemorySegment segment, reader(key)); - return std::make_shared(segment, std::move(eviction_callback)); + return std::make_shared(segment, std::move(eviction_callback)); }; PAIMON_ASSIGN_OR_RAISE(std::shared_ptr cache_value, cache->Get(key, supplier)); - return cache_value->GetSegment(); + auto segment = CacheValue::Get(cache_value); + if (!segment) { + return Status::Invalid("Page cache returned an invalid value type"); + } + return *segment; } void CacheManager::InvalidPage(const std::shared_ptr& key) { diff --git a/src/paimon/common/io/cache/lru_cache_test.cpp b/src/paimon/common/io/cache/lru_cache_test.cpp index 6a9cab5b2..4149516fb 100644 --- a/src/paimon/common/io/cache/lru_cache_test.cpp +++ b/src/paimon/common/io/cache/lru_cache_test.cpp @@ -49,7 +49,7 @@ class LruCacheTest : public ::testing::Test { CacheCallback callback = {}) const { auto segment = MemorySegment::AllocateHeapMemory(size, pool_.get()); std::memset(segment.MutableData(), fill_byte, size); - return std::make_shared(segment, std::move(callback)); + return std::make_shared(segment, std::move(callback)); } auto MakeSupplier(int32_t size, char fill_byte = 0, CacheCallback callback = {}) const { @@ -77,14 +77,16 @@ TEST_F(LruCacheTest, TestGetCacheHitAndMiss) { // First Get: cache miss, supplier should be called ASSERT_OK_AND_ASSIGN(auto value1, cache.Get(key, supplier)); ASSERT_EQ(supplier_call_count, 1); - ASSERT_EQ(value1->GetSegment().Size(), 64); - ASSERT_EQ(value1->GetSegment().Get(0), 'A'); + auto first_segment = CacheValue::Get(value1); + ASSERT_TRUE(first_segment); + ASSERT_EQ(first_segment->Size(), 64); + ASSERT_EQ(first_segment->Get(0), 'A'); ASSERT_EQ(cache.Size(), 1); // Second Get with same key: cache hit, supplier should NOT be called ASSERT_OK_AND_ASSIGN(auto value2, cache.Get(key, supplier)); ASSERT_EQ(supplier_call_count, 1); - ASSERT_EQ(value2->GetSegment().Get(0), 'A'); + ASSERT_EQ(CacheValue::Get(value2)->Get(0), 'A'); ASSERT_EQ(cache.Size(), 1); ASSERT_EQ(cache.GetMaxWeight(), 1024); } @@ -109,21 +111,24 @@ TEST_F(LruCacheTest, TestPutInsertAndUpdate) { // Verify the updated value is returned ASSERT_OK_AND_ASSIGN(auto result, cache.Get(key, MakeSupplier(0))); - ASSERT_EQ(result->GetSegment().Size(), 128); - ASSERT_EQ(result->GetSegment().Get(0), 'B'); + auto segment = CacheValue::Get(result); + ASSERT_TRUE(segment); + ASSERT_EQ(segment->Size(), 128); + ASSERT_EQ(segment->Get(0), 'B'); } TEST_F(LruCacheTest, TestRetainedMemoryAccounting) { LruCache cache(192); auto segment = MemorySegment::AllocateHeapMemory(64, pool_.get()); - auto value = std::make_shared(segment, CacheCallback(), 128); - ASSERT_EQ(64, value->GetSegment().Size()); + auto value = std::make_shared(segment, CacheCallback(), 128); + ASSERT_EQ(64, CacheValue::Get(value)->Size()); ASSERT_EQ(128, value->GetMemoryUsage()); - ASSERT_EQ(64, CacheValue(segment, CacheCallback()).GetMemoryUsage()); - ASSERT_EQ(64, CacheValue(segment, CacheCallback(), 32).GetMemoryUsage()); + ASSERT_EQ(64, MemorySegmentCacheValue(segment, CacheCallback()).GetMemoryUsage()); + ASSERT_EQ(64, MemorySegmentCacheValue(segment, CacheCallback(), 32).GetMemoryUsage()); + ASSERT_EQ(0, MemorySegmentCacheValue(MemorySegment(), CacheCallback()).GetMemoryUsage()); std::vector> evicted; - auto first = std::make_shared( + auto first = std::make_shared( segment, [&](const std::shared_ptr& key) { evicted.push_back(key); }, 128); auto first_key = MakeKey(0); ASSERT_OK(cache.Put(first_key, first)); @@ -139,14 +144,220 @@ TEST_F(LruCacheTest, TestRetainedMemoryAccounting) { LruCache small_cache(100); ASSERT_OK_AND_ASSIGN( std::shared_ptr loaded, - small_cache.Get(MakeKey(0), - [&](const std::shared_ptr&) - -> Result> { return value; })); + small_cache.Get( + MakeKey(0), + [&](const std::shared_ptr&) -> Result> { + return std::shared_ptr(value); + })); ASSERT_EQ(value, loaded); ASSERT_EQ(0, small_cache.Size()); ASSERT_EQ(0, small_cache.GetCurrentWeight()); } +TEST_F(LruCacheTest, TestObjectCacheReuseAndLifetime) { + LruCache cache(128); + auto key = CacheKey::ForKind("test_file", 0, 64, CacheKind::INDEX_READER); + int32_t loads = 0; + std::weak_ptr weak_object; + std::vector> evicted; + auto supplier = [&](const std::shared_ptr&) -> Result> { + ++loads; + auto object = std::make_shared("loaded metadata"); + weak_object = object; + return std::make_shared>( + object, [&](const std::shared_ptr& removed) { evicted.push_back(removed); }, + 128); + }; + ASSERT_OK_AND_ASSIGN(auto first, cache.Get(key, supplier)); + ASSERT_OK_AND_ASSIGN(auto second, cache.Get(key, supplier)); + auto object = CacheValue::Get(first); + ASSERT_TRUE(object); + std::weak_ptr weak_value = first; + ASSERT_EQ(first, second); + ASSERT_EQ(object, CacheValue::Get(second)); + ASSERT_EQ(1, loads); + ASSERT_EQ(128, cache.GetCurrentWeight()); + first.reset(); + second.reset(); + + // Weight-based eviction releases the cache's ownership, but not an active reader's. + ASSERT_OK(cache.Put(MakeKey(1), MakeValue(128))); + ASSERT_EQ((std::vector>{key}), evicted); + ASSERT_FALSE(weak_object.expired()); + ASSERT_FALSE(weak_value.expired()); + ASSERT_EQ("loaded metadata", *object); + object.reset(); + ASSERT_TRUE(weak_object.expired()); + ASSERT_TRUE(weak_value.expired()); + + ASSERT_OK_AND_ASSIGN(auto reloaded, cache.Get(key, supplier)); + ASSERT_EQ(2, loads); + ASSERT_EQ("loaded metadata", *CacheValue::Get(reloaded)); + cache.InvalidateAll(); +} + +TEST_F(LruCacheTest, TestTypedPayloadAccess) { + ASSERT_EQ(nullptr, CacheValue::Get(nullptr)); + ASSERT_EQ(nullptr, CacheValue::Get(nullptr)); + + auto bytes = MakeValue(16, 'A'); + ASSERT_EQ(nullptr, CacheValue::Get(bytes)); + auto object = std::make_shared>( + std::make_shared("metadata"), CacheCallback(), 16); + ASSERT_EQ(nullptr, CacheValue::Get(bytes)); + ASSERT_EQ(nullptr, CacheValue::Get(object)); + ASSERT_EQ(nullptr, CacheValue::Get(object)); + + auto null_object = std::make_shared>(nullptr, CacheCallback(), 1); + ASSERT_EQ(nullptr, CacheValue::Get(null_object)); + auto empty = std::make_shared(MemorySegment(), CacheCallback()); + auto segment = CacheValue::Get(empty); + ASSERT_TRUE(segment); + ASSERT_EQ(0, segment->Size()); + ASSERT_EQ(nullptr, segment->Data()); +} + +TEST_F(LruCacheTest, TestSegmentViewRetainsAliasedOwnerAfterEviction) { + struct OwnedBuffer { + OwnedBuffer() + : value(MemorySegment::WrapView(data.data(), static_cast(data.size())), + CacheCallback()) {} + + std::string data = "borrowed bytes"; + MemorySegmentCacheValue value; + }; + LruCache cache(64); + auto owner = std::make_shared(); + std::weak_ptr weak_owner = owner; + auto value = std::shared_ptr(owner, &owner->value); + ASSERT_OK(cache.Put(MakeKey(0), value)); + auto segment = CacheValue::Get(value); + ASSERT_TRUE(segment); + owner.reset(); + value.reset(); + + ASSERT_OK(cache.Put(MakeKey(1), MakeValue(64))); + ASSERT_EQ(1, cache.Size()); + ASSERT_EQ(64, cache.GetCurrentWeight()); + ASSERT_FALSE(weak_owner.expired()); + ASSERT_EQ(14, segment->Size()); + ASSERT_EQ("borrowed bytes", std::string(segment->Data(), segment->Size())); + segment.reset(); + ASSERT_TRUE(weak_owner.expired()); +} + +TEST_F(LruCacheTest, TestObjectCacheKindsDoNotCollideWithSegments) { + LruCache cache(1024); + const std::vector kinds = {CacheKind::DATA_FILE_FOOTER, CacheKind::INDEX_READER}; + ASSERT_OK(cache.Put(CacheKey::ForKind("file", 0, 64, kinds[0]), MakeValue(64, 'A'))); + for (size_t i = 1; i < kinds.size(); ++i) { + ASSERT_OK(cache.Put(CacheKey::ForKind("file", 0, 64, kinds[i]), + std::make_shared>( + std::make_shared(i), CacheCallback(), sizeof(size_t)))); + } + ASSERT_EQ(kinds.size(), cache.Size()); + for (size_t i = 0; i < kinds.size(); ++i) { + ASSERT_OK_AND_ASSIGN( + auto value, + cache.Get(CacheKey::ForKind("file", 0, 64, kinds[i]), + [](const std::shared_ptr&) -> Result> { + return Status::Invalid("unexpected cache miss"); + })); + if (i == 0) { + ASSERT_EQ('A', CacheValue::Get(value)->Get(0)); + } else { + ASSERT_EQ(i, *CacheValue::Get(value)); + } + } +} + +TEST_F(LruCacheTest, TestObjectMemoryAccounting) { + LruCache cache(64); + auto object = std::make_shared("metadata"); + auto value = std::make_shared>(object, CacheCallback(), 128); + ASSERT_OK_AND_ASSIGN( + auto loaded, + cache.Get(MakeKey(0), + [&](const std::shared_ptr&) -> Result> { + return std::shared_ptr(value); + })); + ASSERT_EQ(object, CacheValue::Get(loaded)); + ASSERT_EQ(128, loaded->GetMemoryUsage()); + ASSERT_EQ(0, cache.Size()); + ASSERT_EQ(0, cache.GetCurrentWeight()); + ASSERT_EQ(1, ObjectCacheValue(object, CacheCallback(), 0).GetMemoryUsage()); + ASSERT_EQ(1, ObjectCacheValue(object, CacheCallback(), -1).GetMemoryUsage()); + auto replacement = std::make_shared(*object); + ASSERT_OK(cache.Put( + MakeKey(1), std::make_shared>(object, CacheCallback(), 32))); + ASSERT_OK(cache.Put(MakeKey(1), std::make_shared>( + replacement, CacheCallback(), 64))); + ASSERT_OK_AND_ASSIGN(auto replaced, cache.Get(MakeKey(1), MakeSupplier(0))); + ASSERT_EQ(replacement, CacheValue::Get(replaced)); + ASSERT_EQ(64, cache.GetCurrentWeight()); +} + +TEST_F(LruCacheTest, TestPolymorphicValueEquality) { + auto bytes = MakeValue(16, 'A'); + auto equal_bytes = MakeValue(16, 'A'); + auto object = std::make_shared("A"); + std::shared_ptr first = + std::make_shared>(object, CacheCallback(), 16); + std::shared_ptr same = + std::make_shared>(object, CacheCallback(), 16); + std::shared_ptr different = std::make_shared>( + std::make_shared(*object), CacheCallback(), 16); + std::shared_ptr other_type = std::make_shared>( + std::make_shared(1), CacheCallback(), 16); + + ASSERT_TRUE(*bytes == *equal_bytes); + ASSERT_TRUE(*first == *same); + ASSERT_FALSE(*first == *different); + ASSERT_FALSE(*first == *other_type); + ASSERT_FALSE(*other_type == *first); + ASSERT_FALSE(*bytes == *first); + ASSERT_FALSE(*first == *bytes); +} + +TEST_F(LruCacheTest, TestReplaceDifferentValueTypes) { + LruCache cache(128); + auto key = MakeKey(0); + std::vector evicted; + auto bytes = MakeValue(64, 'A', [&](const std::shared_ptr& removed) { + ASSERT_EQ(key, removed); + evicted.push_back(1); + }); + auto object = std::make_shared>( + std::make_shared("reader"), + [&](const std::shared_ptr& removed) { + ASSERT_EQ(key, removed); + evicted.push_back(2); + }, + 32); + auto other_type = std::make_shared>( + std::make_shared(7), + [&](const std::shared_ptr& removed) { + ASSERT_EQ(key, removed); + evicted.push_back(3); + }, + 16); + + ASSERT_OK(cache.Put(key, bytes)); + ASSERT_EQ(64, cache.GetCurrentWeight()); + ASSERT_OK(cache.Put(key, object)); + ASSERT_EQ(32, cache.GetCurrentWeight()); + ASSERT_EQ(std::vector{1}, evicted); + ASSERT_OK(cache.Put(key, other_type)); + ASSERT_EQ(16, cache.GetCurrentWeight()); + ASSERT_EQ(1, cache.Size()); + ASSERT_EQ((std::vector{1, 2}), evicted); + ASSERT_OK_AND_ASSIGN(auto cached, cache.Get(key, MakeSupplier(0))); + ASSERT_EQ(other_type, cached); + cache.InvalidateAll(); + ASSERT_EQ(0, cache.GetCurrentWeight()); + ASSERT_EQ((std::vector{1, 2, 3}), evicted); +} + /// Verifies weight-based eviction: when total weight exceeds max, LRU entries are evicted. TEST_F(LruCacheTest, TestWeightBasedEviction) { // Cache can hold at most 200 bytes @@ -175,7 +386,7 @@ TEST_F(LruCacheTest, TestWeightBasedEviction) { }; ASSERT_OK_AND_ASSIGN(auto result, cache.Get(key0, supplier)); ASSERT_EQ(supplier_called, 1); - ASSERT_EQ(result->GetSegment().Get(0), 'X'); + ASSERT_EQ(CacheValue::Get(result)->Get(0), 'X'); } /// Verifies that eviction invokes the CacheCallback on evicted entries. @@ -216,7 +427,7 @@ TEST_F(LruCacheTest, TestLruOrdering) { // Access key0 via Get to move it to front (most recently used) ASSERT_OK_AND_ASSIGN(auto val, cache.Get(key0, MakeSupplier(0))); - ASSERT_EQ(val->GetSegment().Get(0), 'A'); + ASSERT_EQ(CacheValue::Get(val)->Get(0), 'A'); // Insert key2: should evict key1 (now LRU), NOT key0 std::vector evicted; @@ -410,13 +621,28 @@ TEST_F(LruCacheTest, TestForKindSetsKeyKind) { ASSERT_OK_AND_ASSIGN(auto value, cache.Get(key, supplier)); ASSERT_TRUE(supplier_seen_kind); ASSERT_EQ(CacheKind::MANIFEST, key->GetKind()); - ASSERT_EQ('M', value->GetSegment().Get(0)); + ASSERT_EQ('M', CacheValue::Get(value)->Get(0)); auto put_key = CacheKey::ForKind("test_file", 1, 64, CacheKind::MANIFEST); ASSERT_OK(cache.Put(put_key, MakeValue(64, 'P'))); ASSERT_EQ(CacheKind::MANIFEST, put_key->GetKind()); } +TEST_F(LruCacheTest, TestIndexReaderKeys) { + auto key = CacheKey::ForIndexReader("file.index", "embedding"); + auto same = CacheKey::ForIndexReader("file.index", "embedding"); + ASSERT_EQ(CacheKind::INDEX_READER, key->GetKind()); + ASSERT_TRUE(key->Equals(*same)); + ASSERT_TRUE(same->Equals(*key)); + ASSERT_EQ(key->HashCode(), same->HashCode()); + ASSERT_FALSE(key->Equals(*CacheKey::ForIndexReader("other.index", "embedding"))); + ASSERT_FALSE(key->Equals(*CacheKey::ForIndexReader("file.index", "other"))); + + auto uri_with_separator = CacheKey::ForIndexReader("file:index", "column"); + auto column_with_separator = CacheKey::ForIndexReader("file", "index:column"); + ASSERT_FALSE(uri_with_separator->Equals(*column_with_separator)); +} + TEST_F(LruCacheTest, InferredManifestCacheKeysIncludeBucketCountAndSchema) { auto key = SnapshotLiveManifestEntriesCacheKey::ForInferred("table", "main", 1, 4, 0); auto same = SnapshotLiveManifestEntriesCacheKey::ForInferred("table", "main", 1, 4, 0); diff --git a/src/paimon/core/io/file_index_evaluator.cpp b/src/paimon/core/io/file_index_evaluator.cpp index 59c20f477..d33ba7ef7 100644 --- a/src/paimon/core/io/file_index_evaluator.cpp +++ b/src/paimon/core/io/file_index_evaluator.cpp @@ -30,6 +30,7 @@ #include "arrow/type.h" #include "fmt/format.h" #include "fmt/ranges.h" +#include "paimon/cache/cache.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/io/data_file_meta.h" @@ -175,20 +176,20 @@ Result>> OpenSearchReaders( return format_reader->ReadColumnIndex(field_name, &c_schema); } -Result> ExecuteVectorSearch( +Result> GetVectorSearchReader( const std::vector>& readers, - const std::shared_ptr& vector_search) { + const std::string& field_name) { if (readers.empty()) { return Status::NotImplemented( fmt::format("No File Index reader supports vector search for field '{}'", - vector_search->field_name)); + field_name)); } if (readers.size() != 1) { return Status::Invalid( fmt::format("Multiple File Index readers exist for vector search field '{}'", - vector_search->field_name)); + field_name)); } - return readers[0]->VisitVectorSearch(vector_search); + return readers[0]; } Result> ExecuteFullTextSearch( @@ -291,14 +292,47 @@ Result> FileIndexEvaluator::EvaluateVecto const std::shared_ptr& vector_search, const std::shared_ptr& path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, - const std::shared_ptr& pool) { + const std::shared_ptr& pool, const std::shared_ptr& cache) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, ExtractIndexInputStream(/*only_use_embedding_index=*/false, path_factory, file_meta, file_system)); - PAIMON_ASSIGN_OR_RAISE(std::vector> readers, - OpenSearchReaders(data_schema, vector_search->field_name, options, - file_meta, input_stream, pool)); - return ExecuteVectorSearch(readers, vector_search); + auto open_reader = [&]() -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + OpenSearchReaders(data_schema, vector_search->field_name, options, + file_meta, input_stream, pool)); + return GetVectorSearchReader(readers, vector_search->field_name); + }; + std::string uri; + if (cache && input_stream) { + auto uri_result = input_stream->GetUri(); + if (uri_result.ok()) { + uri = std::move(uri_result).value(); + } + } + if (uri.empty()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, open_reader()); + return reader->VisitVectorSearch(vector_search); + } + std::shared_ptr field = data_schema->GetFieldByName(vector_search->field_name); + if (!field) { + return Status::Invalid(fmt::format("Search field '{}' does not exist in data schema", + vector_search->field_name)); + } + auto key = CacheKey::ForIndexReader(uri, vector_search->field_name); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr value, + cache->Get( + key, [&](const std::shared_ptr&) -> Result> { + PAIMON_ASSIGN_OR_RAISE(int64_t length, input_stream->Length()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, open_reader()); + return std::make_shared>(std::move(reader), + nullptr, length); + })); + auto reader = CacheValue::Get(value); + if (!reader) { + return Status::Invalid("Vector index reader cache returned an invalid value"); + } + return reader->VisitVectorSearch(vector_search); } Result> FileIndexEvaluator::EvaluateFullTextSearch( diff --git a/src/paimon/core/io/file_index_evaluator.h b/src/paimon/core/io/file_index_evaluator.h index 82b58ee0f..9a4b50e9a 100644 --- a/src/paimon/core/io/file_index_evaluator.h +++ b/src/paimon/core/io/file_index_evaluator.h @@ -35,6 +35,7 @@ class Schema; } // namespace arrow namespace paimon { +class Cache; class LeafPredicate; class CompoundPredicate; class DataFilePathFactory; @@ -75,7 +76,8 @@ class FileIndexEvaluator { const std::shared_ptr& vector_search, const std::shared_ptr& path_factory, const std::shared_ptr& file_meta, - const std::shared_ptr& file_system, const std::shared_ptr& pool); + const std::shared_ptr& file_system, const std::shared_ptr& pool, + const std::shared_ptr& cache); static Result> EvaluateFullTextSearch( const std::shared_ptr& data_schema, diff --git a/src/paimon/core/io/file_index_evaluator_test.cpp b/src/paimon/core/io/file_index_evaluator_test.cpp index 1a0168efc..05b2f1115 100644 --- a/src/paimon/core/io/file_index_evaluator_test.cpp +++ b/src/paimon/core/io/file_index_evaluator_test.cpp @@ -40,6 +40,7 @@ #include "paimon/core/stats/simple_stats.h" #include "paimon/data/timestamp.h" #include "paimon/defs.h" +#include "paimon/factories/factory_creator.h" #include "paimon/file_index/bitmap_index_result.h" #include "paimon/file_index/file_index_format.h" #include "paimon/file_index/scored_file_index_result.h" @@ -52,6 +53,7 @@ #include "paimon/predicate/vector_search.h" #include "paimon/status.h" #include "paimon/testing/mock/mock_search_file_index.h" +#include "paimon/testing/utils/counting_cache_test_utils.h" #include "paimon/testing/utils/testharness.h" #include "paimon/utils/roaring_bitmap32.h" @@ -488,7 +490,7 @@ class FileIndexSearchEvaluatorTest : public FileIndexEvaluatorTest, if (index_type == SearchIndexType::VECTOR) { ASSERT_NOK_WITH_MSG(FileIndexEvaluator::EvaluateVectorSearch( data_schema_, core_options_.ToMap(), MakeVectorSearch(field), - path_factory_, meta, fs_, pool_), + path_factory_, meta, fs_, pool_, /*cache=*/nullptr), message); } else { ASSERT_NOK_WITH_MSG(FileIndexEvaluator::EvaluateFullTextSearch( @@ -544,7 +546,7 @@ TEST_P(FileIndexSearchEvaluatorTest, TestVectorSearchReadsIndexWithScores) { ASSERT_OK_AND_ASSIGN(std::shared_ptr result, FileIndexEvaluator::EvaluateVectorSearch( data_schema_, core_options_.ToMap(), MakeVectorSearch("embedding"), - path_factory_, meta, fs_, pool_)); + path_factory_, meta, fs_, pool_, /*cache=*/nullptr)); ASSERT_TRUE(result); EXPECT_EQ(RoaringBitmap32::From({1, 3}), result->GetRowPositions()); EXPECT_EQ(std::vector({0.25f, 0.75f}), result->GetScores()); @@ -572,7 +574,7 @@ TEST_P(FileIndexSearchEvaluatorTest, TestEmptyVectorSearchIndex) { ASSERT_OK_AND_ASSIGN(std::shared_ptr vector_result, FileIndexEvaluator::EvaluateVectorSearch( data_schema_, core_options_.ToMap(), MakeVectorSearch("embedding"), - path_factory_, meta, fs_, pool_)); + path_factory_, meta, fs_, pool_, /*cache=*/nullptr)); ASSERT_TRUE(vector_result); EXPECT_TRUE(vector_result->IsEmpty()); EXPECT_TRUE(vector_result->GetScores().empty()); @@ -599,6 +601,292 @@ TEST_P(FileIndexSearchEvaluatorTest, TestInvalidFullTextSearchIndexes) { CheckInvalidSearchIndexes(SearchIndexType::FULL_TEXT); } +namespace { + +struct SearchReaderOpenState { + int32_t count = 0; + bool fail = false; +}; + +class CacheableSearchReader final : public FileIndexReader { + public: + explicit CacheableSearchReader(std::shared_ptr delegate) + : delegate_(std::move(delegate)) {} + + Result> VisitVectorSearch( + const std::shared_ptr& search) override { + if (search->query.empty()) { + return Status::Invalid("Mock vector query cannot be empty"); + } + return delegate_->VisitVectorSearch(search); + } + + private: + std::shared_ptr delegate_; +}; + +class CacheableSearchIndexer final : public FileIndexer { + public: + explicit CacheableSearchIndexer(std::shared_ptr state) + : state_(std::move(state)), delegate_(SearchIndexType::VECTOR) {} + + Result> CreateReader( + ::ArrowSchema* schema, int64_t start, int64_t length, + const std::shared_ptr& input, + const std::shared_ptr& pool) const override { + ++state_->count; + if (state_->fail) { + return Status::IOError("Cannot open mock vector index"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader, + delegate_.CreateReader(schema, start, length, input, pool)); + return std::make_shared(std::move(reader)); + } + + Result> CreateWriter( + ::ArrowSchema* schema, const std::shared_ptr& pool) const override { + return delegate_.CreateWriter(schema, pool); + } + + private: + std::shared_ptr state_; + MockSearchFileIndexer delegate_; +}; + +class CacheableSearchFactory final : public FileIndexerFactory { + public: + explicit CacheableSearchFactory(std::shared_ptr state) + : state_(std::move(state)) {} + + const char* Identifier() const override { + return "mock-cacheable-vector"; + } + + Result> Create( + const std::map&) const override { + return std::make_unique(state_); + } + + private: + std::shared_ptr state_; +}; + +class InvalidSearchReaderCache final : public CountingRoutingCache { + public: + InvalidSearchReaderCache() : CountingRoutingCache(CacheKind::INDEX_READER, 4096) {} + + Result> Get( + const std::shared_ptr&, + std::function>(const std::shared_ptr&)>) + override { + return std::make_shared>(std::make_shared(7), nullptr, + sizeof(int32_t)); + } +}; + +} // namespace + +class FileIndexCacheEvaluatorTest : public FileIndexSearchEvaluatorTest { + protected: + void SetUp() override { + FileIndexSearchEvaluatorTest::SetUp(); + open_state_ = std::make_shared(); + auto factory = std::make_unique(open_state_); + identifier_ = factory->Identifier(); + auto* factories = FactoryCreator::GetInstance(); + factories->Register(identifier_, factory.release()); + cache_ = std::make_shared(CacheKind::INDEX_READER, 1024 * 1024); + } + + void TearDown() override { + auto* factories = FactoryCreator::GetInstance(); + factories->TEST_Unregister(identifier_); + FileIndexSearchEvaluatorTest::TearDown(); + } + + Result> MakeCacheableIndex() const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, + MockSearchFileIndex::MakeIndex("embedding", identifier_, + {{3, 0.75f}, {1, 0.25f}}, pool_)); + return MakeIndex(bytes); + } + + Result> Search( + const std::shared_ptr& meta, + const std::shared_ptr& search) const { + return FileIndexEvaluator::EvaluateVectorSearch(data_schema_, core_options_.ToMap(), search, + path_factory_, meta, fs_, pool_, cache_); + } + + std::shared_ptr open_state_; + std::string identifier_; + std::shared_ptr cache_; +}; + +TEST_P(FileIndexCacheEvaluatorTest, ReusesReaderWithIndependentQueries) { + ASSERT_OK_AND_ASSIGN(auto meta, MakeCacheableIndex()); + ASSERT_OK_AND_ASSIGN(auto first, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(first->GetRowPositions(), RoaringBitmap32::From({1, 3})); + ASSERT_EQ(first->GetScores(), std::vector({0.25f, 0.75f})); + + auto filtered = MakeVectorSearch("embedding"); + filtered->limit = 1; + filtered->query = {5.0f, 6.0f}; + filtered->pre_filter = [](int64_t row_id) { return row_id != 3; }; + ASSERT_OK_AND_ASSIGN(auto second, Search(meta, filtered)); + ASSERT_EQ(second->GetRowPositions(), RoaringBitmap32::From({1})); + ASSERT_EQ(second->GetScores(), std::vector({0.25f})); + + auto invalid = MakeVectorSearch("embedding"); + invalid->query.clear(); + ASSERT_NOK_WITH_MSG(Search(meta, invalid), "Mock vector query cannot be empty"); + ASSERT_OK_AND_ASSIGN(auto last, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(last->GetRowPositions(), first->GetRowPositions()); + ASSERT_EQ(last->GetScores(), first->GetScores()); + ASSERT_EQ(open_state_->count, GetParam() ? 4 : 1); + ASSERT_EQ(cache_->GetCount(), GetParam() ? 0 : 4); + ASSERT_EQ(cache_->SupplierCallCount(), GetParam() ? 0 : 1); + ASSERT_EQ(cache_->Size(), GetParam() ? 0 : 1); +} + +TEST_P(FileIndexCacheEvaluatorTest, DisablesCache) { + ASSERT_OK_AND_ASSIGN(auto meta, MakeCacheableIndex()); + for (int32_t i = 0; i < 2; ++i) { + ASSERT_OK_AND_ASSIGN(auto result, + FileIndexEvaluator::EvaluateVectorSearch( + data_schema_, core_options_.ToMap(), MakeVectorSearch("embedding"), + path_factory_, meta, fs_, pool_, /*cache=*/nullptr)); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32::From({1, 3})); + } + ASSERT_EQ(open_state_->count, 2); + ASSERT_EQ(cache_->GetCount(), 0); + ASSERT_EQ(cache_->Size(), 0); +} + +TEST_P(FileIndexCacheEvaluatorTest, CachesEmptyIndex) { + FileIndexFormat::ColumnIndexes indexes = {{"embedding", {{identifier_, nullptr}}}}; + ASSERT_OK_AND_ASSIGN(auto bytes, MockSearchFileIndex::Serialize(indexes, pool_)); + ASSERT_OK_AND_ASSIGN(auto meta, MakeIndex(bytes)); + for (int32_t i = 0; i < 2; ++i) { + ASSERT_OK_AND_ASSIGN(auto result, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32()); + ASSERT_EQ(result->GetScores(), std::vector()); + } + ASSERT_EQ(open_state_->count, 0); + ASSERT_EQ(cache_->Size(), GetParam() ? 0 : 1); + ASSERT_EQ(cache_->SupplierCallCount(), GetParam() ? 0 : 1); +} + +TEST_P(FileIndexCacheEvaluatorTest, CachesReaderAlongsideUnregisteredIndex) { + ASSERT_OK_AND_ASSIGN(auto payload, MockSearchFileIndex::MakePayload({{1, 0.25f}}, pool_)); + FileIndexFormat::ColumnIndexes indexes = { + {"embedding", {{identifier_, payload}, {"unregistered-index", payload}}}}; + ASSERT_OK_AND_ASSIGN(auto bytes, MockSearchFileIndex::Serialize(indexes, pool_)); + ASSERT_OK_AND_ASSIGN(auto meta, MakeIndex(bytes)); + for (int32_t i = 0; i < 2; ++i) { + ASSERT_OK_AND_ASSIGN(auto result, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32::From({1})); + ASSERT_EQ(result->GetScores(), std::vector({0.25f})); + } + ASSERT_EQ(open_state_->count, GetParam() ? 2 : 1); + ASSERT_EQ(cache_->Size(), GetParam() ? 0 : 1); + ASSERT_EQ(cache_->SupplierCallCount(), GetParam() ? 0 : 1); +} + +TEST_P(FileIndexCacheEvaluatorTest, RejectsMultipleReaders) { + ASSERT_OK_AND_ASSIGN(auto payload, MockSearchFileIndex::MakePayload({{1, 0.25f}}, pool_)); + FileIndexFormat::ColumnIndexes indexes = { + {"embedding", {{identifier_, payload}, {"another-empty-index", nullptr}}}}; + ASSERT_OK_AND_ASSIGN(auto bytes, MockSearchFileIndex::Serialize(indexes, pool_)); + ASSERT_OK_AND_ASSIGN(auto meta, MakeIndex(bytes)); + ASSERT_NOK_WITH_MSG(Search(meta, MakeVectorSearch("embedding")), + "Multiple File Index readers exist for vector search field 'embedding'"); + ASSERT_EQ(cache_->Size(), 0); + ASSERT_EQ(cache_->SupplierCallCount(), GetParam() ? 0 : 1); +} + +TEST_P(FileIndexCacheEvaluatorTest, ReloadsAfterEvictionAndDoesNotCacheOpenFailures) { + ASSERT_OK_AND_ASSIGN( + auto bytes, + MockSearchFileIndex::MakeIndex("embedding", identifier_, {{3, 0.75f}, {1, 0.25f}}, pool_)); + ASSERT_OK_AND_ASSIGN(auto meta, MakeIndex(bytes)); + cache_ = std::make_shared(CacheKind::INDEX_READER, bytes->size() - 1); + for (int32_t i = 0; i < 2; ++i) { + ASSERT_OK_AND_ASSIGN(auto result, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32::From({1, 3})); + } + ASSERT_EQ(open_state_->count, 2); + ASSERT_EQ(cache_->Size(), 0); + + cache_ = std::make_shared(CacheKind::INDEX_READER, bytes->size()); + open_state_->fail = true; + ASSERT_NOK_WITH_MSG(Search(meta, MakeVectorSearch("embedding")), + "Cannot open mock vector index"); + ASSERT_EQ(cache_->Size(), 0); + open_state_->fail = false; + ASSERT_OK_AND_ASSIGN(auto result, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32::From({1, 3})); + ASSERT_EQ(open_state_->count, 4); + ASSERT_EQ(cache_->Size(), GetParam() ? 0 : 1); + cache_->InvalidateAll(); + ASSERT_OK_AND_ASSIGN(result, Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(result->GetRowPositions(), RoaringBitmap32::From({1, 3})); + ASSERT_EQ(open_state_->count, 5); +} + +TEST_P(FileIndexCacheEvaluatorTest, KeysReadersByFileAndField) { + ASSERT_OK_AND_ASSIGN(auto payload, MockSearchFileIndex::MakePayload({{1, 0.25f}}, pool_)); + FileIndexFormat::ColumnIndexes indexes; + indexes["embedding"][identifier_] = payload; + indexes["title"][identifier_] = payload; + ASSERT_OK_AND_ASSIGN(auto bytes, MockSearchFileIndex::Serialize(indexes, pool_)); + ASSERT_OK_AND_ASSIGN(auto meta, MakeIndex(bytes)); + ASSERT_OK(Search(meta, MakeVectorSearch("embedding"))); + ASSERT_OK(Search(meta, MakeVectorSearch("title"))); + ASSERT_EQ(open_state_->count, 2); + auto options = core_options_.ToMap(); + options["file-index." + identifier_ + ".columns"] = "embedding"; + options["file-index." + identifier_ + ".embedding.variant"] = "second"; + ASSERT_OK(FileIndexEvaluator::EvaluateVectorSearch(data_schema_, options, + MakeVectorSearch("embedding"), path_factory_, + meta, fs_, pool_, cache_)); + ASSERT_EQ(open_state_->count, GetParam() ? 3 : 2); + auto schema = arrow::schema({arrow::field("embedding", arrow::list(arrow::float64()))}); + ASSERT_OK(FileIndexEvaluator::EvaluateVectorSearch(schema, core_options_.ToMap(), + MakeVectorSearch("embedding"), path_factory_, + meta, fs_, pool_, cache_)); + ASSERT_EQ(open_state_->count, GetParam() ? 4 : 2); + if (!GetParam()) { + const std::string copy_name = "another-data.orc.index"; + ASSERT_OK(fs_->WriteFile(path_factory_->ToAlignedPath(copy_name, meta), + std::string(bytes->data(), bytes->size()), /*overwrite=*/true)); + meta->extra_files = {copy_name}; + } + ASSERT_OK(Search(meta, MakeVectorSearch("embedding"))); + ASSERT_EQ(open_state_->count, GetParam() ? 5 : 3); + ASSERT_EQ(cache_->Size(), GetParam() ? 0 : 3); +} + +TEST_P(FileIndexCacheEvaluatorTest, RejectsWrongCacheValue) { + ASSERT_OK_AND_ASSIGN(auto meta, MakeCacheableIndex()); + cache_ = std::make_shared(); + auto result = Search(meta, MakeVectorSearch("embedding")); + if (GetParam()) { + ASSERT_OK(result); + ASSERT_EQ(result.value()->GetRowPositions(), RoaringBitmap32::From({1, 3})); + ASSERT_EQ(open_state_->count, 1); + } else { + ASSERT_FALSE(result.ok()); + ASSERT_TRUE(result.status().IsInvalid()); + ASSERT_EQ(open_state_->count, 0); + } +} + +INSTANTIATE_TEST_SUITE_P(IndexStorage, FileIndexCacheEvaluatorTest, ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Embedded" : "External"; + }); + INSTANTIATE_TEST_SUITE_P(IndexStorage, FileIndexSearchEvaluatorTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { return info.param ? "Embedded" : "External"; diff --git a/src/paimon/core/manifest/index_manifest_file_handler_test.cpp b/src/paimon/core/manifest/index_manifest_file_handler_test.cpp index e0ce213a3..fd2ea4535 100644 --- a/src/paimon/core/manifest/index_manifest_file_handler_test.cpp +++ b/src/paimon/core/manifest/index_manifest_file_handler_test.cpp @@ -128,9 +128,10 @@ TEST_F(IndexManifestFileHandlerTest, DecodedCacheReusesIpcAcrossReaders) { return Status::Invalid("index manifest was not cached"); })); ASSERT_TRUE(cached); - const auto& segment = cached->GetSegment(); - auto buffer = std::make_shared(reinterpret_cast(segment.Data()), - segment.Size()); + auto segment = CacheValue::Get(cached); + ASSERT_TRUE(segment); + auto buffer = std::make_shared(reinterpret_cast(segment->Data()), + segment->Size()); auto ipc = arrow::ipc::RecordBatchStreamReader::Open( std::make_shared(buffer)); ASSERT_TRUE(ipc.ok()) << ipc.status().ToString(); diff --git a/src/paimon/core/manifest/manifest_list_test.cpp b/src/paimon/core/manifest/manifest_list_test.cpp index 8653eaf88..b28b772b1 100644 --- a/src/paimon/core/manifest/manifest_list_test.cpp +++ b/src/paimon/core/manifest/manifest_list_test.cpp @@ -153,9 +153,10 @@ TEST_F(ManifestListTest, DecodedCacheReusesIpcAcrossReaders) { return Status::Invalid("manifest list was not cached"); })); ASSERT_TRUE(cached); - const auto& segment = cached->GetSegment(); - auto buffer = std::make_shared(reinterpret_cast(segment.Data()), - segment.Size()); + auto segment = CacheValue::Get(cached); + ASSERT_TRUE(segment); + auto buffer = std::make_shared(reinterpret_cast(segment->Data()), + segment->Size()); // The existing cache key contains IPC, not source-format bytes. auto ipc = arrow::ipc::RecordBatchStreamReader::Open( std::make_shared(buffer)); diff --git a/src/paimon/core/manifest/manifest_row_range_test.cpp b/src/paimon/core/manifest/manifest_row_range_test.cpp index db2dd75cd..a5827277d 100644 --- a/src/paimon/core/manifest/manifest_row_range_test.cpp +++ b/src/paimon/core/manifest/manifest_row_range_test.cpp @@ -552,7 +552,7 @@ TEST_F(RowRangeManifestFileTest, CacheChargesRetainedArrowCapacity) { return Status::Invalid("manifest was not cached"); })); ASSERT_TRUE(cached); - const int64_t logical_size = cached->GetSegment().Size(); + const int64_t logical_size = CacheValue::Get(cached)->Size(); ASSERT_GT(cached->GetMemoryUsage(), logical_size); ASSERT_EQ(pool->CurrentUsage(), cached->GetMemoryUsage()); cached.reset(); diff --git a/src/paimon/core/operation/append_only_file_store_scan_test.cpp b/src/paimon/core/operation/append_only_file_store_scan_test.cpp index b8dfbfced..e3d5cf1b8 100644 --- a/src/paimon/core/operation/append_only_file_store_scan_test.cpp +++ b/src/paimon/core/operation/append_only_file_store_scan_test.cpp @@ -1052,7 +1052,7 @@ TEST(AppendOnlyFileStoreScanTest, SnapshotCacheRebuildsCorruptValues) { auto scan = BuildScan(table_path, cache, /*bucket=*/0, predicate, lazy_decode); scan->WithSnapshot(snapshot); ASSERT_OK(cache->Put(scan->CreateSnapshotLiveManifestEntriesCacheKey(0), - std::make_shared(segment, CacheCallback()))); + std::make_shared(segment, CacheCallback()))); for (int32_t attempt = 0; attempt < 2; ++attempt) { ASSERT_OK_AND_ASSIGN(auto plan, scan->CreatePlan()); ASSERT_EQ(SortedFileNames(plan->Files()), expected_names); diff --git a/src/paimon/core/operation/file_store_scan.cpp b/src/paimon/core/operation/file_store_scan.cpp index ca2de308a..45f32405d 100644 --- a/src/paimon/core/operation/file_store_scan.cpp +++ b/src/paimon/core/operation/file_store_scan.cpp @@ -411,9 +411,10 @@ Result FileStoreScan::ReadManifestEntriesWithCache( [](const std::shared_ptr&) -> Result> { return std::shared_ptr(); }); - if (value.ok() && value.value()) { + auto segment = value.ok() ? CacheValue::Get(value.value()) : nullptr; + if (segment) { Result> selected = - SnapshotLiveManifestEntries::ReadSnapshot(value.value()->GetSegment(), snapshot.Id(), + SnapshotLiveManifestEntries::ReadSnapshot(*segment, snapshot.Id(), max_snapshots, row_filter, pool_); if (predicate_failed) { return selected.status(); @@ -428,7 +429,7 @@ Result FileStoreScan::ReadManifestEntriesWithCache( // Reuse the same cache value to retain complete historical snapshots on a miss. // Query-filtered entries must never be written back into the cache. Result deserialized = - SnapshotLiveManifestEntries::Deserialize(value.value()->GetSegment(), max_snapshots, + SnapshotLiveManifestEntries::Deserialize(*segment, max_snapshots, pool_); if (deserialized.ok()) { cached_entries = std::move(deserialized.value()); @@ -499,7 +500,7 @@ Status FileStoreScan::StoreSnapshotLiveManifestEntries( return Status::OK(); } auto cache_value = - std::make_shared(MemorySegment::Wrap(bytes_result.value()), CacheCallback()); + std::make_shared(MemorySegment::Wrap(bytes_result.value()), CacheCallback()); Status status = core_options_.GetCache()->Put(CreateSnapshotLiveManifestEntriesCacheKey(bucket), cache_value); return status.ok() ? status : Status::OK(); diff --git a/src/paimon/core/operation/raw_file_split_read.cpp b/src/paimon/core/operation/raw_file_split_read.cpp index 58a19fc6b..3675e3501 100644 --- a/src/paimon/core/operation/raw_file_split_read.cpp +++ b/src/paimon/core/operation/raw_file_split_read.cpp @@ -314,7 +314,8 @@ Result> RawFileSplitRead::EvaluateVectorS } return FileIndexEvaluator::EvaluateVectorSearch(data_schema, file_index_options, search_with_dv, data_file_path_factory, file, - options_.GetFileSystem(), pool_); + options_.GetFileSystem(), pool_, + options_.GetCache()); } Result> RawFileSplitRead::EvaluateFullTextSearchWithDv( diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 77f92109c..2a62ca9ea 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -34,11 +34,13 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.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/deletionvectors/bitmap64_deletion_vector.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" +#include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/manifest/file_source.h" #include "paimon/core/operation/internal_read_context.h" #include "paimon/core/schema/schema_manager.h" @@ -51,6 +53,7 @@ #include "paimon/factories/factory_creator.h" #include "paimon/format/file_format.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/bytes.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/full_text_search.h" #include "paimon/predicate/predicate_builder.h" @@ -60,10 +63,56 @@ #include "paimon/table/source/data_split.h" #include "paimon/testing/mock/mock_search_file_index.h" #include "paimon/testing/utils/binary_row_generator.h" +#include "paimon/testing/utils/counting_cache_test_utils.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +class CacheTrackingFileIndexer : public FileIndexer { + public: + explicit CacheTrackingFileIndexer(std::shared_ptr open_count) + : open_count_(std::move(open_count)), delegate_(SearchIndexType::VECTOR) {} + + Result> CreateReader( + ::ArrowSchema* schema, int64_t start, int64_t length, + const std::shared_ptr& input, + const std::shared_ptr& pool) const override { + ++*open_count_; + return delegate_.CreateReader(schema, start, length, input, pool); + } + + Result> CreateWriter( + ::ArrowSchema* schema, const std::shared_ptr& pool) const override { + return delegate_.CreateWriter(schema, pool); + } + + private: + std::shared_ptr open_count_; + MockSearchFileIndexer delegate_; +}; + +class CacheTrackingFileIndexFactory : public FileIndexerFactory { + public: + explicit CacheTrackingFileIndexFactory(std::shared_ptr open_count) + : open_count_(std::move(open_count)) {} + + const char* Identifier() const override { + return "lumina"; + } + + Result> Create( + const std::map&) const override { + return std::make_unique(open_count_); + } + + private: + std::shared_ptr open_count_; +}; + +} // namespace + class RawFileSplitReadTest : public ::testing::Test { protected: std::vector> PrepareDataSplits() const { @@ -456,13 +505,28 @@ TEST_F(RawFileSplitReadTest, TestCreateReaderWithNonPartitionWithReserveSequence } TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { + auto open_count = std::make_shared(0); + auto factory = std::make_unique(open_count); + const std::string index_type = factory->Identifier(); + auto* factories = FactoryCreator::GetInstance(); + auto original_factory = factories->factories_.extract(index_type); + factories->Register(index_type, factory.release()); + ScopeGuard unregister_factory([&]() { + factories->TEST_Unregister(index_type); + if (!original_factory.empty()) { + factories->factories_.insert(std::move(original_factory)); + } + }); + auto cache = std::make_shared(CacheKind::INDEX_READER, 1024 * 1024); std::string path = paimon::test::GetDataDir() + "/orc/multi_partition_append_table.db/multi_partition_append_table"; std::shared_ptr vector_search = std::make_shared( "f3", /*limit=*/2, std::vector{1.0f}, nullptr, nullptr, VectorSearch::DistanceType::EUCLIDEAN, std::map{}); ReadContextBuilder context_builder(path); - context_builder.SetReadFieldNames({"_INDEX_SCORE", "f0"}).SetVectorSearch(vector_search); + context_builder.SetReadFieldNames({"_INDEX_SCORE", "f0"}) + .SetVectorSearch(vector_search) + .WithCache(cache); ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); SchemaManager schema_manager(std::make_shared(), path); ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); @@ -474,9 +538,8 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { auto data_split = std::dynamic_pointer_cast(data_splits.front()); ASSERT_TRUE(data_split); ASSERT_OK_AND_ASSIGN( - data_split->DataFiles().front()->embedded_index, - MockSearchFileIndex::MakeIndex("f3", MockSearchFileIndex::VECTOR_IDENTIFIER, - {{0, 0.25f}, {2, 0.75f}}, pool_)); + std::shared_ptr index_bytes, + MockSearchFileIndex::MakeIndex("f3", index_type, {{0, 0.25f}, {2, 0.75f}}, pool_)); const CoreOptions& core_options = internal_context->GetCoreOptions(); std::shared_ptr data_schema = @@ -493,6 +556,21 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { core_options.DataFilePrefix(), core_options.LegacyPartitionNameEnabled(), external_paths, global_index_external_path, core_options.IndexFileInDataFileDir(), pool_)); + auto index_dir = UniqueTestDirectory::Create(); + auto file_meta = data_split->DataFiles().front(); + auto fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN( + auto data_path_factory, + path_factory->CreateDataFilePathFactory(data_split->Partition(), data_split->Bucket())); + std::string data_bytes; + ASSERT_OK(fs->ReadFile(data_path_factory->ToPath(file_meta), &data_bytes)); + file_meta->external_path = index_dir->Str() + "/" + file_meta->file_name; + ASSERT_OK(fs->WriteFile(file_meta->external_path.value(), data_bytes, /*overwrite=*/true)); + const std::string index_name = file_meta->file_name + ".index"; + ASSERT_OK(fs->WriteFile(data_path_factory->ToAlignedPath(index_name, file_meta), + std::string(index_bytes->data(), index_bytes->size()), + /*overwrite=*/true)); + file_meta->extra_files = {index_name}; ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(/*thread_count=*/2)); auto split_read = std::make_unique(path_factory, std::move(internal_context), @@ -509,6 +587,15 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { {R"([[0, 0.25, "Bob"], [0, 0.75, "Tony"]])"}, &expected) .ok()); ASSERT_TRUE(result->Equals(expected)); + ASSERT_EQ(*open_count, 1); + ASSERT_EQ(cache->SupplierCallCount(CacheKind::INDEX_READER), 1); + ASSERT_EQ(cache->Size(), 1); + ASSERT_OK_AND_ASSIGN(reader, split_read->CreateReader(data_split)); + ASSERT_OK_AND_ASSIGN(result, ReadResultCollector::CollectResult(std::move(reader))); + ASSERT_TRUE(result->Equals(expected)); + ASSERT_EQ(*open_count, 1); + ASSERT_EQ(cache->GetCount(CacheKind::INDEX_READER), 2); + ASSERT_EQ(cache->SupplierCallCount(CacheKind::INDEX_READER), 1); ReadContextBuilder without_score_builder(path); without_score_builder.SetReadFieldNames({"f0"}).SetVectorSearch(vector_search); @@ -530,6 +617,8 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { {R"([[0, "Bob"], [0, "Tony"]])"}, &without_score_expected) .ok()); ASSERT_TRUE(without_score_result->Equals(without_score_expected)); + ASSERT_EQ(*open_count, 2); + ASSERT_EQ(cache->GetCount(CacheKind::INDEX_READER), 2); } TEST_F(RawFileSplitReadTest, TestVectorSearchFiltersDeletionVectorBeforeTopK) { diff --git a/src/paimon/core/utils/objects_file.h b/src/paimon/core/utils/objects_file.h index 2a99e98d1..36ba10e05 100644 --- a/src/paimon/core/utils/objects_file.h +++ b/src/paimon/core/utils/objects_file.h @@ -302,7 +302,7 @@ Result> ObjectsFile::SerializeArrowBatches( CacheCallback(), data->capacity()) {} std::shared_ptr pool; std::shared_ptr buffer; - CacheValue value; + MemorySegmentCacheValue value; }; auto owner = std::make_shared(arrow_pool_, buffer); auto* value = &owner->value; @@ -324,7 +324,6 @@ Status ObjectsFile::ReadArrowBatches( auto cached = cache_->Get(key, [&](const std::shared_ptr&) { return SerializeArrowBatches(file_name, file_size, &batches, &read_status); }); - const bool usable = cached.ok() && cached.value() && cached.value()->GetSegment().Data(); if (read_status) { // Publish the query-independent cache before invoking user filters. A filter error // does not invalidate the cache or affect other readers of this file. @@ -334,12 +333,12 @@ Status ObjectsFile::ReadArrowBatches( } return Status::OK(); } - if (!usable) { + auto segment = cached.ok() ? CacheValue::Get(cached.value()) : nullptr; + if (!segment || !segment->Data()) { return ReadUncachedArrowBatches(file_name, file_size, consumer, prepare_reader); } - const auto& segment = cached.value()->GetSegment(); - auto buffer = std::make_shared(reinterpret_cast(segment.Data()), - segment.Size()); + auto buffer = std::make_shared(reinterpret_cast(segment->Data()), + segment->Size()); auto input = std::make_shared(buffer); auto read_options = arrow::ipc::IpcReadOptions::Defaults(); read_options.memory_pool = arrow_pool_.get(); diff --git a/src/paimon/format/parquet/parquet_input_stream.h b/src/paimon/format/parquet/parquet_input_stream.h index fbf45f244..e5650cfd9 100644 --- a/src/paimon/format/parquet/parquet_input_stream.h +++ b/src/paimon/format/parquet/parquet_input_stream.h @@ -105,15 +105,16 @@ class ParquetInputStream : public ArrowInputStreamAdapter { if (size != nbytes) { return Status::IOError("Short read of Parquet page index"); } - return std::make_shared(segment, CacheCallback()); + return std::make_shared(segment, CacheCallback()); }); if (!value.ok()) { return ToArrowStatus(value.status()); } - if (!value.value() || value.value()->GetSegment().Size() != nbytes) { + auto segment = CacheValue::Get(value.value()); + if (!segment || segment->Size() != nbytes) { return arrow::Status::IOError("Invalid Parquet page-index cache value"); } - std::memcpy(out, value.value()->GetSegment().Data(), nbytes); + std::memcpy(out, segment->Data(), nbytes); return nbytes; } diff --git a/src/paimon/format/parquet/parquet_reader_builder.h b/src/paimon/format/parquet/parquet_reader_builder.h index fad961ace..13dd9633e 100644 --- a/src/paimon/format/parquet/parquet_reader_builder.h +++ b/src/paimon/format/parquet/parquet_reader_builder.h @@ -165,15 +165,16 @@ class ParquetReaderBuilder : public ReaderBuilder { PAIMON_ASSIGN_OR_RAISE( MemorySegment segment, SerializeParquetMetadataFooter(input_stream, reader_properties, arrow_pool)); - return std::make_shared(segment, CacheCallback()); + return std::make_shared(segment, CacheCallback()); }; PAIMON_ASSIGN_OR_RAISE(std::shared_ptr cache_value, cache_->Get(cache_key, supplier)); - if (cache_value == nullptr) { - return Status::Invalid("Parquet metadata cache returned nullptr value"); + auto segment = CacheValue::Get(cache_value); + if (!segment) { + return Status::Invalid("Parquet metadata cache returned an invalid value type"); } - return ParseParquetMetadataFooter(cache_value->GetSegment(), reader_properties); + return ParseParquetMetadataFooter(*segment, reader_properties); } int32_t batch_size_ = -1; diff --git a/src/paimon/indexer/lumina/CMakeLists.txt b/src/paimon/indexer/lumina/CMakeLists.txt index 58ba4daad..280956e7c 100644 --- a/src/paimon/indexer/lumina/CMakeLists.txt +++ b/src/paimon/indexer/lumina/CMakeLists.txt @@ -57,6 +57,7 @@ if(PAIMON_ENABLE_LUMINA) test_utils_static ${PAIMON_LUMINA_INDEX_STATIC_LINK_LIBS} ${PAIMON_LOCAL_FILE_SYSTEM_STATIC_LINK_LIBS} + paimon_mock_file_format_static lumina::interface ${GTEST_LINK_TOOLCHAIN}) endif() diff --git a/src/paimon/indexer/lumina/lumina_file_index_test.cpp b/src/paimon/indexer/lumina/lumina_file_index_test.cpp index a3ddd8e66..ee7de80bf 100644 --- a/src/paimon/indexer/lumina/lumina_file_index_test.cpp +++ b/src/paimon/indexer/lumina/lumina_file_index_test.cpp @@ -18,6 +18,8 @@ #include "paimon/indexer/lumina/lumina_file_index.h" +#include +#include #include #include #include @@ -31,10 +33,16 @@ #include "paimon/common/io/byte_array_output_stream.h" #include "paimon/common/io/offset_output_stream.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/io/file_index_evaluator.h" +#include "paimon/file_index/file_index_format.h" #include "paimon/io/byte_array_input_stream.h" #include "paimon/memory/memory_pool.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/predicate/vector_search.h" +#include "paimon/testing/mock/mock_file_format_factory.h" +#include "paimon/testing/utils/counting_cache_test_utils.h" #include "paimon/testing/utils/testharness.h" #include "paimon/utils/roaring_bitmap32.h" @@ -302,4 +310,252 @@ TEST(LuminaFileIndexTest, RoundTripWithTagPredicate) { EXPECT_FLOAT_EQ(4.0f, result->GetScores()[0]); } +namespace { + +class LuminaFileIndexCacheTest : public ::testing::Test { + protected: + void SetUp() override { + dir_ = paimon::test::UniqueTestDirectory::Create("local"); + cache_ = std::make_shared(CacheKind::INDEX_READER, + 16 * 1024 * 1024); + path_factory_ = std::make_shared(); + ASSERT_OK(path_factory_->Init(dir_->Str(), /*format_identifier=*/"orc", + /*data_file_prefix=*/"data-", nullptr)); + const std::map index_options = { + {"index.dimension", "4"}, + {"index.type", "bruteforce"}, + {"distance.metric", "l2"}, + {"encoding.type", "rawf32"}, + }; + for (const auto& [key, value] : index_options) { + options_["file-index.lumina.embedding." + key] = value; + } + options_["file-index.lumina.columns"] = "embedding"; + LuminaFileIndexer indexer(index_options); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr writer, + indexer.CreateWriter(&c_schema, pool_)); + ::ArrowArray c_batch; + ASSERT_TRUE(arrow::ExportArray(*CreateVectors(), &c_batch).ok()); + ASSERT_OK(writer->AddBatch(&c_batch)); + auto output = + std::make_shared(std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr format_writer, + FileIndexFormat::CreateWriter(output, FileIndexFormat::V_2, pool_)); + ASSERT_OK(format_writer->WriteIndex("embedding", "lumina", [&](OutputStream* payload) { + return writer->WriteTo(payload); + })); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK_AND_ASSIGN(container_, output->Finish(pool_.get())); + ASSERT_OK_AND_ASSIGN(meta_, + DataFileMeta::ForAppend( + "data-vectors.orc", /*file_size=*/0, /*row_count=*/3, + SimpleStats::EmptyStats(), /*min_sequence_number=*/0, + /*max_sequence_number=*/2, /*schema_id=*/0, FileSource::Append(), + /*value_stats_cols=*/std::nullopt, /*external_path=*/std::nullopt, + /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt)); + ASSERT_OK(WriteIndex("vectors.index")); + } + + Status WriteIndex(const std::string& name) { + meta_->extra_files = {name}; + return dir_->GetFileSystem()->WriteFile(path_factory_->ToAlignedPath(name, meta_), + std::string(container_->data(), container_->size()), + /*overwrite=*/false); + } + + Result> Search( + const std::shared_ptr& query) const { + return FileIndexEvaluator::EvaluateVectorSearch( + schema_, options_, query, path_factory_, meta_, dir_->GetFileSystem(), pool_, cache_); + } + + static std::shared_ptr Query(bool filter) { + auto query = std::make_shared( + "embedding", 1, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, nullptr, nullptr, + VectorSearch::DistanceType::EUCLIDEAN, std::map{}); + if (filter) { + query->pre_filter = [](int64_t row) { return row == 0; }; + } + return query; + } + + static void CheckResult(const std::shared_ptr& result, bool filter) { + ASSERT_TRUE(result); + ASSERT_EQ(RoaringBitmap32::From({filter ? 0 : 2}), result->GetRowPositions()); + ASSERT_EQ(std::vector({filter ? 4.0f : 0.0f}), result->GetScores()); + } + + std::unique_ptr dir_; + std::shared_ptr pool_ = GetDefaultPool(); + std::shared_ptr cache_; + std::shared_ptr path_factory_; + std::shared_ptr meta_; + std::shared_ptr container_; + std::shared_ptr schema_ = + arrow::schema({arrow::field("embedding", arrow::list(arrow::float32()))}); + std::map options_ = { + {"file.format", paimon::test::MockFileFormatFactory::IDENTIFIER}, + {"manifest.format", paimon::test::MockFileFormatFactory::IDENTIFIER}}; +}; + +TEST_F(LuminaFileIndexCacheTest, ReusesSearcherAndKeepsConcurrentQueriesIndependent) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(1, cache_->SupplierCallCount()); + ASSERT_EQ(CacheKind::INDEX_READER, cache_->LastKind()); + ASSERT_EQ(1, cache_->Size()); + + std::promise start; + std::shared_future ready = start.get_future().share(); + std::vector> searches; + constexpr int32_t kThreads = 8; + constexpr int32_t kQueriesPerThread = 32; + for (int32_t i = 0; i < kThreads; ++i) { + searches.push_back(std::async(std::launch::async, [&, i, ready] { + ready.wait(); + bool filter = i % 2 == 0; + for (int32_t j = 0; j < kQueriesPerThread; ++j) { + auto query = Query(filter); + query->options["lumina.search.parallel_number"] = "2"; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr found, Search(query)); + if (!(found->GetRowPositions() == RoaringBitmap32::From({filter ? 0 : 2})) || + found->GetScores() != std::vector({filter ? 4.0f : 0.0f})) { + return Status::Invalid("Concurrent queries returned mixed results"); + } + } + return Status::OK(); + })); + } + start.set_value(); + for (auto& search : searches) { + ASSERT_OK(search.get()); + } + ASSERT_EQ(1, cache_->SupplierCallCount()); + ASSERT_EQ(1 + kThreads * kQueriesPerThread, cache_->GetCount(CacheKind::INDEX_READER)); + ASSERT_EQ(1, cache_->Size()); +} + +TEST_F(LuminaFileIndexCacheTest, QueryErrorsKeepSearcherAndInFlightEvictionIsSafe) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + auto invalid_query = Query(false); + invalid_query->query.clear(); + ASSERT_NOK(Search(invalid_query)); + ASSERT_EQ(1, cache_->Size()); + ASSERT_OK_AND_ASSIGN(result, Search(Query(true))); + CheckResult(result, true); + ASSERT_EQ(1, cache_->SupplierCallCount()); + + auto evicting_query = Query(false); + evicting_query->pre_filter = [&](int64_t row) { + cache_->InvalidateAll(); + return row == 0; + }; + ASSERT_OK_AND_ASSIGN(result, Search(evicting_query)); + CheckResult(result, true); + ASSERT_EQ(0, cache_->Size()); + ASSERT_EQ(1, cache_->SupplierCallCount()); + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(2, cache_->SupplierCallCount()); + ASSERT_EQ(1, cache_->Size()); +} + +TEST_F(LuminaFileIndexCacheTest, UsesExternallyPopulatedReader) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr input, + dir_->GetFileSystem()->Open(path_factory_->ToAlignedPath("vectors.index", meta_))); + ASSERT_OK_AND_ASSIGN(std::string uri, input->GetUri()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr format_reader, + FileIndexFormat::CreateReader(input, pool_, options_)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema_, &c_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::vector> readers, + format_reader->ReadColumnIndex("embedding", &c_schema)); + ASSERT_EQ(1, readers.size()); + auto key = CacheKey::ForIndexReader(uri, "embedding"); + ASSERT_OK(cache_->Put(key, std::make_shared>( + readers[0], nullptr, container_->size()))); + readers.clear(); + format_reader.reset(); + input.reset(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + ASSERT_OK_AND_ASSIGN(result, Search(Query(true))); + CheckResult(result, true); + ASSERT_EQ(2, cache_->GetCount(CacheKind::INDEX_READER)); + ASSERT_EQ(0, cache_->SupplierCallCount()); + ASSERT_EQ(1, cache_->Size()); + + cache_->Invalidate(CacheKey::ForIndexReader(uri, "embedding")); + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(1, cache_->SupplierCallCount()); +} + +TEST_F(LuminaFileIndexCacheTest, ReusesReaderAcrossOptionsAndSeparatesFiles) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(1, cache_->SupplierCallCount()); + + options_["file-index.lumina.embedding.distance.metric"] = "inner_product"; + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(1, cache_->SupplierCallCount()); + ASSERT_EQ(1, cache_->Size()); + + options_["file-index.lumina.embedding.distance.metric"] = "l2"; + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(1, cache_->SupplierCallCount()); + + ASSERT_OK(WriteIndex("other.index")); + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_EQ(2, cache_->SupplierCallCount()); + ASSERT_EQ(2, cache_->Size()); +} + +TEST_F(LuminaFileIndexCacheTest, OversizedObjectsRemainUsableWithoutRetention) { + cache_ = std::make_shared(CacheKind::INDEX_READER, + container_->size() - 1); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + ASSERT_OK_AND_ASSIGN(result, Search(Query(true))); + CheckResult(result, true); + ASSERT_EQ(2, cache_->SupplierCallCount()); + ASSERT_EQ(0, cache_->Size()); + + cache_ = std::make_shared(CacheKind::INDEX_READER, + container_->size()); + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); + ASSERT_OK_AND_ASSIGN(result, Search(Query(true))); + CheckResult(result, true); + ASSERT_EQ(1, cache_->SupplierCallCount()); + ASSERT_EQ(1, cache_->Size()); +} + +TEST_F(LuminaFileIndexCacheTest, BypassesEmbeddedIndexesAndDisabledCache) { + meta_->embedded_index = container_; + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, Search(Query(false))); + CheckResult(result, false); + ASSERT_OK_AND_ASSIGN(result, Search(Query(true))); + CheckResult(result, true); + ASSERT_EQ(0, cache_->GetCount()); + ASSERT_EQ(0, cache_->SupplierCallCount()); + ASSERT_EQ(0, cache_->Size()); + + meta_->embedded_index.reset(); + cache_.reset(); + ASSERT_OK_AND_ASSIGN(result, Search(Query(false))); + CheckResult(result, false); +} + +} // namespace + } // namespace paimon::lumina::test diff --git a/src/paimon/indexer/lumina/lumina_index_searcher.cpp b/src/paimon/indexer/lumina/lumina_index_searcher.cpp index 78c025556..d8e5acb11 100644 --- a/src/paimon/indexer/lumina/lumina_index_searcher.cpp +++ b/src/paimon/indexer/lumina/lumina_index_searcher.cpp @@ -79,6 +79,7 @@ Result<::lumina::api::LuminaSearcher::SearchResult> LuminaIndexSearcher::Search( PAIMON_ASSIGN_OR_RAISE(::lumina::api::SearchOptions search_options, LuminaIndexOptions::CreateSearchOptions(*vector_search, index_info_)); ::lumina::api::Query query(vector_search->query.data(), vector_search->query.size()); + LuminaMemoryPool session_pool(pool_->GetPaimonPool()); if (vector_search->predicate) { if (!searcher_with_tag_) { @@ -89,7 +90,7 @@ Result<::lumina::api::LuminaSearcher::SearchResult> LuminaIndexSearcher::Search( if (!vector_search->pre_filter) { PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( ::lumina::api::LuminaSearcher::SearchResult search_result, - searcher_with_tag_->SearchWithTag(query, tag_filter, search_options, *pool_)); + searcher_with_tag_->SearchWithTag(query, tag_filter, search_options, session_pool)); return std::move(search_result); } auto filter = [pre_filter = vector_search->pre_filter]( @@ -97,14 +98,14 @@ Result<::lumina::api::LuminaSearcher::SearchResult> LuminaIndexSearcher::Search( PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( ::lumina::api::LuminaSearcher::SearchResult search_result, searcher_with_tag_->SearchWithTagAndFilter(query, tag_filter, filter, search_options, - *pool_)); + session_pool)); return std::move(search_result); } if (!vector_search->pre_filter) { PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( ::lumina::api::LuminaSearcher::SearchResult search_result, - searcher_->Search(query, search_options, *pool_)); + searcher_->Search(query, search_options, session_pool)); return std::move(search_result); } auto filter = [pre_filter = vector_search->pre_filter](::lumina::core::vector_id_t id) -> bool { @@ -112,7 +113,7 @@ Result<::lumina::api::LuminaSearcher::SearchResult> LuminaIndexSearcher::Search( }; PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( ::lumina::api::LuminaSearcher::SearchResult search_result, - searcher_with_filter_->SearchWithFilter(query, filter, search_options, *pool_)); + searcher_with_filter_->SearchWithFilter(query, filter, search_options, session_pool)); return std::move(search_result); }