This is an automated email from the ASF dual-hosted git repository.

yiguolei pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/branch-4.2 by this push:
     new cadac8d3180 branch-4.2: [fix](mow) Account for pinned PK roots in 
segment cache eviction (#69275)
cadac8d3180 is described below

commit cadac8d318016fe135a848ca72c284cc33773d8d
Author: hui lai <[email protected]>
AuthorDate: Tue Oct 20 10:06:10 2026 +0800

    branch-4.2: [fix](mow) Account for pinned PK roots in segment cache 
eviction (#69275)
    
    ### What problem does this PR solve?
    
    Backport https://github.com/apache/doris/pull/68304 to `branch-4.2`.
    
    SegmentCache could retain large decompressed PK index roots while
    charging only estimated metadata. Account for those pinned roots in
    eviction weight at insertion and after deferred index initialization,
    without counting the same physical pages twice in memory tracking. Grow
    charges under the LRU shard lock, guard updates against replacement,
    keep the callback's cache ownership weak, and release partially loaded
    index readers on failure. Existing eager rowset index loading and
    index-before-bloom-filter lookup order are preserved.
    
    Adapt the Segment test fixture to branch-4.2's existing
    `SegmentWriterOptions` and `build_segment` helper. Retain all 7 LRU and
    6 Segment test cases from the merged source change.
    
    ### Release note
    
    Account for pinned PK index roots in SegmentCache eviction to reduce
    persistent index memory retention.
    
    ### Check List (For Author)
    
    - Test: Passed `git diff --check`, changed-line/new-file clang-format
    16.0.5 checks, and local include-path verification. Compilation and test
    execution intentionally skipped as requested. This release branch does
    not contain `build-support/check-build-hygiene.sh`.
    - Behavior changed: Yes; loaded PK roots affect SegmentCache eviction
    weight. Evicted segments remain usable through active shared pointers;
    more frequent eviction may increase index reloads.
    - Does this need documentation: No.
---
 be/src/storage/index/primary_key_index.cpp         |  12 +-
 be/src/storage/segment/segment.cpp                 |   6 +-
 be/src/storage/segment/segment.h                   |  16 ++
 be/src/storage/segment/segment_loader.cpp          |  21 ++-
 be/src/util/lru_cache.cpp                          |  47 +++++
 be/src/util/lru_cache.h                            |   7 +
 be/test/storage/cache/lru_cache_test.cpp           | 175 +++++++++++++++++++
 .../segment/segment_primary_key_cache_test.cpp     | 190 +++++++++++++++++++++
 8 files changed, 468 insertions(+), 6 deletions(-)

diff --git a/be/src/storage/index/primary_key_index.cpp 
b/be/src/storage/index/primary_key_index.cpp
index 28fa377e4b4..080e19d503b 100644
--- a/be/src/storage/index/primary_key_index.cpp
+++ b/be/src/storage/index/primary_key_index.cpp
@@ -112,12 +112,16 @@ Status 
PrimaryKeyIndexReader::parse_index(io::FileReaderSPtr file_reader,
                                           OlapReaderStatistics* 
pk_index_load_stats,
                                           const io::IOContext* source_io_ctx) {
     // parse primary key index
-    _index_reader.reset(new segment_v2::IndexedColumnReader(file_reader, 
meta.primary_key_index()));
-    _index_reader->set_is_pk_index(true);
+    auto index_reader = 
std::make_unique<segment_v2::IndexedColumnReader>(file_reader,
+                                                                          
meta.primary_key_index());
+    index_reader->set_is_pk_index(true);
     auto io_ctx = create_index_io_context(source_io_ctx, pk_index_load_stats);
-    
RETURN_IF_ERROR(_index_reader->load(!config::disable_pk_storage_page_cache, 
false,
-                                        pk_index_load_stats, &io_ctx));
+    RETURN_IF_ERROR(index_reader->load(!config::disable_pk_storage_page_cache, 
false,
+                                       pk_index_load_stats, &io_ctx));
 
+    // A failed initialization must not leave partially loaded root pages 
pinned in
+    // a cached segment without a corresponding eviction charge.
+    _index_reader = std::move(index_reader);
     _index_parsed = true;
     return Status::OK();
 }
diff --git a/be/src/storage/segment/segment.cpp 
b/be/src/storage/segment/segment.cpp
index 0913e80f364..6258e2110e5 100644
--- a/be/src/storage/segment/segment.cpp
+++ b/be/src/storage/segment/segment.cpp
@@ -730,7 +730,11 @@ Status Segment::load_index(OlapReaderStatistics* stats, 
const io::IOContext* sou
             _pk_index_reader = std::make_unique<PrimaryKeyIndexReader>();
             RETURN_IF_ERROR(_pk_index_reader->parse_index(_file_reader, 
*_pk_index_meta, stats,
                                                           source_io_ctx));
-            // _meta_mem_usage += _pk_index_reader->get_memory_size();
+            _pk_index_cache_bytes.store(_pk_index_reader->get_memory_size(),
+                                        std::memory_order_relaxed);
+            if (_cache_charge_callback) {
+                _cache_charge_callback(cache_charge());
+            }
             return Status::OK();
         } else {
             // read and parse short key index page
diff --git a/be/src/storage/segment/segment.h b/be/src/storage/segment/segment.h
index e236e22da9e..41199a31379 100644
--- a/be/src/storage/segment/segment.h
+++ b/be/src/storage/segment/segment.h
@@ -22,7 +22,9 @@
 #include <gen_cpp/segment_v2.pb.h>
 #include <glog/logging.h>
 
+#include <atomic>
 #include <cstdint>
+#include <functional>
 #include <map>
 #include <memory> // for unique_ptr
 #include <string>
@@ -177,6 +179,18 @@ public:
     // another method `get_metadata_size` not include the column reader, only 
the segment object itself.
     int64_t meta_mem_usage() const { return _meta_mem_usage; }
 
+    // PK pages are tracked by PKIndexPageCache, but their pinned bytes must 
also
+    // count toward SegmentCache eviction. Do not add these bytes to metadata 
tracking.
+    size_t cache_charge() const {
+        return _meta_mem_usage + 
_pk_index_cache_bytes.load(std::memory_order_relaxed);
+    }
+
+    // Install before publishing this segment in SegmentCache; immutable 
afterwards.
+    void set_cache_charge_callback(std::function<void(size_t)> callback) {
+        DCHECK(!_cache_charge_callback);
+        _cache_charge_callback = std::move(callback);
+    }
+
     // Get the inner file column's data type.
     // When `read_options` is provided, the decision (e.g. flat-leaf vs 
hierarchical) can depend
     // on the reader type and tablet schema; when it is nullptr, we treat it 
as a query reader.
@@ -278,6 +292,8 @@ private:
     // The memory consumed by querying is tracked in segment iterator.
     int64_t _meta_mem_usage;
     int64_t _tracked_meta_mem_usage = 0;
+    std::atomic<size_t> _pk_index_cache_bytes {0};
+    std::function<void(size_t)> _cache_charge_callback;
 
     RowsetId _rowset_id;
     TabletSchemaSPtr _tablet_schema;
diff --git a/be/src/storage/segment/segment_loader.cpp 
b/be/src/storage/segment/segment_loader.cpp
index 49d704b68eb..81d9a20bc3a 100644
--- a/be/src/storage/segment/segment_loader.cpp
+++ b/be/src/storage/segment/segment_loader.cpp
@@ -42,8 +42,27 @@ bool SegmentCache::lookup(const SegmentCache::CacheKey& key, 
SegmentCacheHandle*
 
 void SegmentCache::insert(const SegmentCache::CacheKey& key, 
SegmentCache::CacheValue& value,
                           SegmentCacheHandle* handle) {
+    // Keep the cache weakly referenced: the cache owns the segment. The raw 
segment
+    // pointer is only compared while its own initialization invokes the 
callback.
+    value.segment->set_cache_charge_callback([cache = 
std::weak_ptr<Cache>(_cache),
+                                              encoded_key = key.encode(),
+                                              segment = 
value.segment.get()](size_t charge) {
+        auto owner = cache.lock();
+        if (!owner) {
+            return; // The Segment may outlive the cache during shutdown.
+        }
+        auto* entry = owner->lookup(encoded_key);
+        if (entry == nullptr) {
+            return; // Already evicted.
+        }
+        // Another instance of this segment may have replaced the cached entry.
+        if (static_cast<CacheValue*>(owner->value(entry))->segment.get() == 
segment) {
+            owner->update_charge(entry, charge);
+        }
+        owner->release(entry);
+    });
     auto* lru_handle =
-            LRUCachePolicy::insert(key.encode(), &value, 
value.segment->meta_mem_usage(),
+            LRUCachePolicy::insert(key.encode(), &value, 
value.segment->cache_charge(),
                                    value.segment->meta_mem_usage(), 
CachePriority::NORMAL);
     handle->push_segment(this, lru_handle);
 }
diff --git a/be/src/util/lru_cache.cpp b/be/src/util/lru_cache.cpp
index ff7b868ccc6..2076fe208af 100644
--- a/be/src/util/lru_cache.cpp
+++ b/be/src/util/lru_cache.cpp
@@ -336,6 +336,48 @@ Cache::Handle* LRUCache::lookup(const CacheKey& key, 
uint32_t hash) {
     return reinterpret_cast<Cache::Handle*>(e);
 }
 
+void LRUCache::update_charge(Cache::Handle* handle, size_t charge) {
+    auto* e = reinterpret_cast<LRUHandle*>(handle);
+    LRUHandle* to_remove_head = nullptr;
+    {
+        std::lock_guard l(_mutex);
+        // Erase/replacement may have removed this entry while its handle was 
held.
+        if (!e->in_cache || charge <= e->charge) {
+            return;
+        }
+        DCHECK_GT(e->refs, 1);
+        size_t delta = charge - e->charge;
+        e->charge = charge;
+        e->total_size += delta;
+        _usage += delta;
+        if (e->total_size > _capacity) {
+            // This entry cannot fit even after evicting every other entry. 
Drop
+            // cache ownership now, while outstanding handles keep it alive.
+            bool removed = _table.remove(e);
+            DCHECK(removed);
+            e->in_cache = false;
+            _unref(e);
+            _usage -= e->total_size;
+        }
+        // Growing an existing entry does not need another element slot. The
+        // insertion eviction helpers also check count >= limit, so only call
+        // them when the updated usage exceeds the byte capacity.
+        if (_usage > _capacity) {
+            if (_cache_value_check_timestamp) {
+                _evict_from_lru_with_time(0, &to_remove_head);
+            } else {
+                _evict_from_lru(0, &to_remove_head);
+            }
+        }
+    }
+    // Cache values can release other caches; never destroy them under the 
shard lock.
+    while (to_remove_head != nullptr) {
+        LRUHandle* next = to_remove_head->next;
+        to_remove_head->free();
+        to_remove_head = next;
+    }
+}
+
 void LRUCache::release(Cache::Handle* handle) {
     if (handle == nullptr) {
         return;
@@ -770,6 +812,11 @@ void ShardedLRUCache::release(Handle* handle) {
     _shards[_shard(h->hash)]->release(handle);
 }
 
+void ShardedLRUCache::update_charge(Handle* handle, size_t charge) {
+    auto* h = reinterpret_cast<LRUHandle*>(handle);
+    _shards[_shard(h->hash)]->update_charge(handle, charge);
+}
+
 void ShardedLRUCache::erase(const CacheKey& key) {
     const uint32_t hash = _hash_slice(key);
     _shards[_shard(hash)]->erase(key, hash);
diff --git a/be/src/util/lru_cache.h b/be/src/util/lru_cache.h
index 1e77b9a0422..53f17cd48ae 100644
--- a/be/src/util/lru_cache.h
+++ b/be/src/util/lru_cache.h
@@ -194,6 +194,10 @@ public:
     // REQUIRES: handle must have been returned by a method on *this.
     virtual void release(Handle* handle) = 0;
 
+    // Grow an entry's eviction charge after lazy initialization, without 
changing its
+    // memory tracker. The caller must hold a live handle. Stale smaller 
charges are ignored.
+    virtual void update_charge(Handle* handle, size_t charge) = 0;
+
     // Return the value encapsulated in a handle returned by a
     // successful lookup().
     // REQUIRES: handle must not have been released yet.
@@ -349,6 +353,7 @@ public:
                           CachePriority priority = CachePriority::NORMAL);
     Cache::Handle* lookup(const CacheKey& key, uint32_t hash);
     void release(Cache::Handle* handle);
+    void update_charge(Cache::Handle* handle, size_t charge);
     void erase(const CacheKey& key, uint32_t hash);
     PrunedInfo prune();
     PrunedInfo prune_if(CachePrunePredicate pred, bool lazy_mode = false);
@@ -421,6 +426,7 @@ public:
                    CachePriority priority = CachePriority::NORMAL) override;
     Handle* lookup(const CacheKey& key) override;
     void release(Handle* handle) override;
+    void update_charge(Handle* handle, size_t charge) override;
     void erase(const CacheKey& key) override;
     void* value(Handle* handle) override;
     uint64_t new_id() override;
@@ -485,6 +491,7 @@ public:
                    CachePriority priority = CachePriority::NORMAL) override;
     Handle* lookup(const CacheKey& key) override { return nullptr; };
     void release(Handle* handle) override;
+    void update_charge(Handle* handle, size_t charge) override {}
     void erase(const CacheKey& key) override {};
     void* value(Handle* handle) override;
     uint64_t new_id() override { return 0; };
diff --git a/be/test/storage/cache/lru_cache_test.cpp 
b/be/test/storage/cache/lru_cache_test.cpp
index a31d532a5a4..5c248838ad6 100644
--- a/be/test/storage/cache/lru_cache_test.cpp
+++ b/be/test/storage/cache/lru_cache_test.cpp
@@ -21,6 +21,7 @@
 #include <gtest/gtest-message.h>
 #include <gtest/gtest-test-part.h>
 
+#include <future>
 #include <iosfwd>
 #include <vector>
 
@@ -297,6 +298,180 @@ TEST_F(CacheTest, EvictionPolicyWithDurable) {
     EXPECT_EQ(201, Lookup(200));
 }
 
+TEST_F(CacheTest, UpdateChargeDoesNotDoubleTrackMemory) {
+    init_size_cache(1024 * 1024);
+    const CacheKey key("growing");
+    auto* handle = cache()->insert(key, new CacheValue(EncodeValue(1)), 100, 
100);
+    const auto usage = cache()->get_usage();
+    const auto tracked = cache()->mem_consumption();
+    cache()->_cache->update_charge(handle, 500);
+    EXPECT_EQ(cache()->get_usage(), usage + 400);
+    EXPECT_EQ(cache()->mem_consumption(), tracked);
+    // Repeated initialization and older snapshots cannot undo the larger 
charge.
+    cache()->_cache->update_charge(handle, 500);
+    cache()->_cache->update_charge(handle, 200);
+    EXPECT_EQ(cache()->get_usage(), usage + 400);
+    cache()->_cache->update_charge(handle, 2 * 1024 * 1024);
+    EXPECT_EQ(cache()->get_usage(), 0);
+    // Detaching an oversized entry does not release its live handle's memory.
+    EXPECT_EQ(cache()->mem_consumption(), tracked);
+    cache()->release(handle);
+    cache()->prune_all(true);
+    EXPECT_EQ(cache()->get_usage(), 0);
+    EXPECT_EQ(cache()->mem_consumption(), 0);
+}
+
+TEST_F(CacheTest, UpdateChargeEvictsIdleEntriesWhenEntryStillFits) {
+    LRUCache cache(LRUCacheType::SIZE);
+    cache.set_capacity(1024);
+    CacheKey idle("idle"), durable("durable"), growing("growing");
+    cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 100));
+    cache.release(
+            cache.insert(durable, 3, new CacheValue(EncodeValue(3)), 100, 
CachePriority::DURABLE));
+    auto* handle = cache.insert(growing, 2, new CacheValue(EncodeValue(2)), 
100);
+    cache.update_charge(handle, 700);
+    EXPECT_EQ(cache.lookup(idle, 1), nullptr);
+    EXPECT_LE(cache.get_usage(), 1024);
+    auto* retained = cache.lookup(durable, 3);
+    ASSERT_NE(retained, nullptr);
+    cache.release(retained);
+    cache.release(handle);
+    EXPECT_EQ(cache.get_element_count(), 2);
+}
+
+TEST_F(CacheTest, UpdateChargeAtElementLimitOnlyEvictsAboveByteCapacity) {
+    for (bool check_timestamp : {false, true}) {
+        LRUCache cache(LRUCacheType::SIZE, true);
+        cache.set_capacity(1024);
+        cache.set_element_count_capacity(2);
+        cache.set_cache_value_time_extractor([](const void* value) -> int64_t {
+            return DecodeValue(static_cast<const CacheValue*>(value)->value);
+        });
+        cache.set_cache_value_check_timestamp(check_timestamp);
+        CacheKey idle("idle"), growing("growing");
+        cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 
100));
+        auto* handle = cache.insert(growing, 2, new 
CacheValue(EncodeValue(2)), 100);
+        ASSERT_EQ(cache.get_element_count(), 2);
+        const auto usage = cache.get_usage();
+
+        cache.update_charge(handle, 200);
+        EXPECT_EQ(cache.get_usage(), usage + 100);
+        EXPECT_EQ(cache.get_element_count(), 2);
+        auto* retained = cache.lookup(idle, 1);
+        ASSERT_NE(retained, nullptr);
+        cache.release(retained);
+
+        // Exactly filling byte capacity must also preserve both entries.
+        const auto full_charge = 100 + cache.get_capacity() - usage;
+        cache.update_charge(handle, full_charge);
+        EXPECT_EQ(cache.get_usage(), cache.get_capacity());
+        EXPECT_EQ(cache.get_element_count(), 2);
+
+        // Crossing byte capacity still evicts the idle entry normally.
+        cache.update_charge(handle, full_charge + 1);
+        EXPECT_EQ(cache.lookup(idle, 1), nullptr);
+        EXPECT_EQ(cache.get_element_count(), 1);
+        EXPECT_LE(cache.get_usage(), cache.get_capacity());
+        cache.release(handle);
+        auto* growing_handle = cache.lookup(growing, 2);
+        ASSERT_NE(growing_handle, nullptr);
+        cache.release(growing_handle);
+    }
+}
+
+TEST_F(CacheTest, UpdateChargeOversizedEntryPreservesOtherEntries) {
+    for (bool check_timestamp : {false, true}) {
+        LRUCache cache(LRUCacheType::SIZE, true);
+        cache.set_capacity(1024);
+        cache.set_cache_value_time_extractor([](const void* value) -> int64_t {
+            return DecodeValue(static_cast<const CacheValue*>(value)->value);
+        });
+        cache.set_cache_value_check_timestamp(check_timestamp);
+        CacheKey idle("idle"), durable("durable"), growing("growing");
+        cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 
100));
+        cache.release(cache.insert(durable, 3, new CacheValue(EncodeValue(3)), 
100,
+                                   CachePriority::DURABLE));
+        const auto usage = cache.get_usage();
+        auto* handle = cache.insert(growing, 2, new 
CacheValue(EncodeValue(2)), 100);
+        cache.update_charge(handle, 2048);
+        EXPECT_EQ(cache.get_usage(), usage);
+        EXPECT_EQ(cache.get_element_count(), 2);
+        EXPECT_EQ(cache.lookup(growing, 2), nullptr);
+        // The detached entry remains usable through the caller's live handle.
+        
EXPECT_EQ(DecodeValue(static_cast<CacheValue*>(reinterpret_cast<LRUHandle*>(handle)->value)
+                                      ->value),
+                  2);
+        cache.release(handle);
+        auto* idle_handle = cache.lookup(idle, 1);
+        ASSERT_NE(idle_handle, nullptr);
+        cache.release(idle_handle);
+        auto* durable_handle = cache.lookup(durable, 3);
+        ASSERT_NE(durable_handle, nullptr);
+        cache.release(durable_handle);
+        EXPECT_EQ(cache.get_usage(), usage);
+    }
+}
+
+TEST_F(CacheTest, UpdateChargeIgnoresReplacedEntry) {
+    LRUCache cache(LRUCacheType::SIZE);
+    cache.set_capacity(4096);
+    CacheKey key("replaced");
+    auto* old = cache.insert(key, 1, new CacheValue(EncodeValue(1)), 100);
+    auto* current = cache.insert(key, 1, new CacheValue(EncodeValue(2)), 200);
+    const auto usage = cache.get_usage();
+    cache.update_charge(old, 8192);
+    EXPECT_EQ(cache.get_usage(), usage);
+    cache.release(old);
+    cache.release(current);
+    EXPECT_EQ(cache.get_usage(), usage);
+}
+
+TEST_F(CacheTest, UpdateChargeConcurrentGrowth) {
+    LRUCache cache(LRUCacheType::SIZE);
+    cache.set_capacity(1024 * 1024);
+    CacheKey key("concurrent");
+    auto* handle = cache.insert(key, 1, new CacheValue(EncodeValue(1)), 100);
+    const auto usage = cache.get_usage();
+    std::vector<std::future<void>> updates;
+    for (size_t charge = 200; charge <= 900; charge += 100) {
+        auto* held = cache.lookup(key, 1);
+        ASSERT_NE(held, nullptr);
+        updates.emplace_back(std::async(std::launch::async, [&, held, charge] {
+            cache.update_charge(held, charge);
+            cache.release(held);
+        }));
+    }
+    for (auto& update : updates) {
+        update.get();
+    }
+    EXPECT_EQ(cache.get_usage(), usage + 800);
+    cache.release(handle);
+}
+
+TEST_F(CacheTest, UpdateChargePreservesOtherOutstandingHandles) {
+    LRUCache cache(LRUCacheType::SIZE);
+    cache.set_capacity(1024);
+    CacheKey key("shared");
+    auto* first = cache.insert(key, 1, new CacheValueWithKey(1, 
EncodeValue(1)), 100);
+    auto* second = cache.lookup(key, 1);
+    ASSERT_NE(second, nullptr);
+    cache.update_charge(first, 2048);
+    EXPECT_EQ(cache.get_usage(), 0);
+    EXPECT_EQ(cache.lookup(key, 1), nullptr);
+    cache.release(first);
+    EXPECT_TRUE(_deleted_keys.empty());
+    EXPECT_EQ(
+            
DecodeValue(static_cast<CacheValueWithKey*>(reinterpret_cast<LRUHandle*>(second)->value)
+                                ->value),
+            1);
+    cache.update_charge(second, 4096);
+    EXPECT_EQ(cache.get_usage(), 0);
+    cache.release(second);
+    ASSERT_EQ(_deleted_keys.size(), 1);
+    EXPECT_EQ(_deleted_keys.front(), 1);
+    EXPECT_EQ(cache.get_usage(), 0);
+}
+
 TEST_F(CacheTest, Usage) {
     LRUCache cache(LRUCacheType::SIZE);
     cache.set_capacity(1040);
diff --git a/be/test/storage/segment/segment_primary_key_cache_test.cpp 
b/be/test/storage/segment/segment_primary_key_cache_test.cpp
new file mode 100644
index 00000000000..977036d2c5f
--- /dev/null
+++ b/be/test/storage/segment/segment_primary_key_cache_test.cpp
@@ -0,0 +1,190 @@
+// 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 <gtest/gtest.h>
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/field.h"
+#include "io/fs/local_file_system.h"
+#include "runtime/exec_env.h"
+#include "storage/index/primary_key_index.h"
+#include "storage/segment/segment.h"
+#include "storage/segment/segment_loader.h"
+#include "storage/segment/segment_writer.h"
+#include "storage/storage_engine.h"
+#include "storage/tablet/tablet_schema_helper.h"
+
+namespace doris {
+
+TabletSchemaSPtr create_schema(const std::vector<TabletColumnPtr>& columns, 
KeysType keys_type);
+using Generator = std::function<void(size_t rid, int cid, Field& field)>;
+void build_segment(SegmentWriterOptions opts, TabletSchemaSPtr build_schema, 
size_t segment_id,
+                   TabletSchemaSPtr query_schema, size_t nrows, Generator 
generator,
+                   std::shared_ptr<Segment>* res, std::string segment_dir);
+
+class SegmentPrimaryKeyCacheTest : public testing::Test {
+protected:
+    void SetUp() override {
+        auto fs = io::global_local_filesystem();
+        ASSERT_TRUE(fs->delete_directory(_dir).ok());
+        ASSERT_TRUE(fs->create_directory(_dir).ok());
+        ExecEnv::GetInstance()->set_storage_engine(
+                std::make_unique<StorageEngine>(EngineOptions {}));
+        _schema = create_schema(
+                {create_varchar_key(0, false),
+                 create_int_value(1, 
FieldAggregationMethod::OLAP_FIELD_AGGREGATION_REPLACE,
+                                  false)},
+                UNIQUE_KEYS);
+        SegmentWriterOptions opts;
+        opts.enable_unique_key_merge_on_write = true;
+        auto generator = [](size_t rid, int cid, Field& field) {
+            if (cid == 0) {
+                // Each key exceeds the normal PK data-page target, producing 
a large value index.
+                field = Field::create_field<TYPE_VARCHAR>(std::string(50000, 
'a' + rid));
+            } else {
+                field = 
Field::create_field<TYPE_INT>(static_cast<int32_t>(rid));
+            }
+        };
+        build_segment(opts, _schema, 0, _schema, 8, generator, &_segment, 
_dir);
+        ASSERT_NE(_segment, nullptr);
+        ASSERT_EQ(_segment->_pk_index_reader, nullptr);
+        ASSERT_FALSE(_segment->_load_index_once.has_called());
+        ASSERT_FALSE(_segment->_load_pk_bf_once.has_called());
+        ASSERT_FALSE(_segment->_pk_index_meta->primary_key_index()
+                             .value_index_meta()
+                             .is_root_data_page());
+        _present_key = _segment->min_key();
+
+        // Find a definite BF miss without initializing the segment's own 
reader. Do not assume
+        // that an arbitrary absent key is rejected, because bloom filters 
allow false positives.
+        PrimaryKeyIndexReader probe;
+        ASSERT_TRUE(
+                probe.parse_bf(_segment->file_reader(), 
*_segment->_pk_index_meta, nullptr).ok());
+        for (int i = 0; i < 10000; ++i) {
+            auto candidate = _present_key + std::to_string(i);
+            if (!probe.check_present(Slice(candidate))) {
+                _missing_key = std::move(candidate);
+                break;
+            }
+        }
+        ASSERT_FALSE(_missing_key.empty());
+    }
+
+    void TearDown() override {
+        _segment.reset();
+        ExecEnv::GetInstance()->set_storage_engine(nullptr);
+        
EXPECT_TRUE(io::global_local_filesystem()->delete_directory(_dir).ok());
+    }
+
+    Status lookup(const std::string& key, RowLocation* location) {
+        return _segment->lookup_row_key(Slice(key), _schema.get(), false, 
false, location, nullptr);
+    }
+
+    const std::string _dir = "./ut_dir/segment_primary_key_cache_test";
+    TabletSchemaSPtr _schema;
+    std::shared_ptr<segment_v2::Segment> _segment;
+    std::string _present_key;
+    std::string _missing_key;
+};
+
+TEST_F(SegmentPrimaryKeyCacheTest, CachedRootPagesIncreaseChargeOnlyOnce) {
+    SegmentCache cache(1024 * 1024 * 1024, 10000);
+    SegmentCacheHandle handle;
+    const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id());
+    cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle);
+    const auto usage = cache.get_usage();
+    const auto tracked = cache.mem_consumption();
+    RowLocation location;
+    ASSERT_TRUE(lookup(_missing_key, 
&location).is<ErrorCode::KEY_NOT_FOUND>());
+    // Keep the original index-first lookup: even a definite BF miss loads the 
roots.
+    const auto roots = _segment->_pk_index_reader->get_memory_size();
+    EXPECT_GT(roots, 300000);
+    EXPECT_EQ(cache.get_usage(), usage + roots);
+    EXPECT_EQ(cache.mem_consumption(), tracked);
+    ASSERT_TRUE(lookup(_present_key, &location).ok());
+    ASSERT_TRUE(_segment->load_index(nullptr).ok());
+    EXPECT_EQ(cache.get_usage(), usage + roots);
+}
+
+TEST_F(SegmentPrimaryKeyCacheTest, LazyRootLoadingEvictsOversizedSegment) {
+    // SegmentCache has 64 shards. Metadata fits; the large PK roots do not.
+    SegmentCache cache((_segment->cache_charge() + 65536) * 64, 10000);
+    SegmentCacheHandle handle;
+    const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id());
+    cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle);
+    ASSERT_EQ(cache.get_element_count(), 1);
+    ASSERT_TRUE(_segment->load_index(nullptr).ok());
+    EXPECT_EQ(cache.get_element_count(), 0);
+    EXPECT_EQ(cache.get_usage(), 0);
+    // Eviction drops cache ownership, not the active operation's shared_ptr.
+    RowLocation location;
+    ASSERT_TRUE(lookup(_present_key, &location).ok());
+    EXPECT_EQ(location.row_id, 0);
+}
+
+TEST_F(SegmentPrimaryKeyCacheTest, EagerRootLoadingChargesOnInsert) {
+    const auto capacity = (_segment->cache_charge() + 65536) * 64;
+    ASSERT_TRUE(_segment->load_index(nullptr).ok());
+    SegmentCache cache(capacity, 10000);
+    SegmentCacheHandle handle;
+    const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id());
+    cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle);
+    EXPECT_EQ(cache.get_usage(), 0);
+    EXPECT_EQ(cache.get_element_count(), 0);
+    EXPECT_EQ(handle.get_segments().front(), _segment);
+}
+
+TEST_F(SegmentPrimaryKeyCacheTest, OldSegmentDoesNotChargeItsReplacement) {
+    SegmentCache cache(1024 * 1024 * 1024, 10000);
+    const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id());
+    SegmentCacheHandle old_handle, new_handle;
+    cache.insert(key, *new SegmentCache::CacheValue(_segment), &old_handle);
+    auto replacement = std::make_shared<segment_v2::Segment>(_segment->id(), 
_segment->rowset_id(),
+                                                             _schema, 
InvertedIndexFileInfo {});
+    cache.insert(key, *new SegmentCache::CacheValue(replacement), &new_handle);
+    const auto usage = cache.get_usage();
+    ASSERT_TRUE(_segment->load_index(nullptr).ok());
+    EXPECT_EQ(cache.get_usage(), usage);
+    EXPECT_EQ(cache.get_element_count(), 1);
+}
+
+TEST_F(SegmentPrimaryKeyCacheTest, SegmentCanOutliveCache) {
+    {
+        SegmentCache cache(1024 * 1024 * 1024, 10000);
+        SegmentCacheHandle handle;
+        const SegmentCache::CacheKey key(_segment->rowset_id(), 
_segment->id());
+        cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle);
+    }
+    ASSERT_TRUE(_segment->load_index(nullptr).ok());
+    RowLocation location;
+    ASSERT_TRUE(lookup(_present_key, &location).ok());
+}
+
+TEST_F(SegmentPrimaryKeyCacheTest, FailedIndexLoadDoesNotPinPartialRoots) {
+    auto meta = *_segment->_pk_index_meta;
+    
meta.mutable_primary_key_index()->mutable_value_index_meta()->mutable_root_page()->set_offset(
+            uint64_t {1} << 50);
+    PrimaryKeyIndexReader reader;
+    EXPECT_FALSE(reader.parse_index(_segment->file_reader(), meta, 
nullptr).ok());
+    EXPECT_EQ(reader._index_reader, nullptr);
+}
+
+} // namespace doris


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to