From 893dc73118bc33ebd16061bf7893dda0152b12ff Mon Sep 17 00:00:00 2001 From: yanghua Date: Sun, 26 Jul 2026 20:53:12 +0800 Subject: [PATCH] feat(index): per-query index cache hit/miss stats Surface per-page index cache hit and miss counts alongside the existing `iops` / `parts_loaded` / `index_comparisons` metrics, so users can tell whether a scan warmed from the in-memory index cache or paid the storage cost. Wiring: - Extend `MetricsCollector` with `record_index_cache_hit/miss` (default no-op) and mirror on `LocalMetricsCollector`, `IndexMetrics`, and `FtsIndexMetrics` (wrapper forwarding). - Add matching `Count` fields to `IndexMetrics` so the two metrics show up automatically in every `ExecutionPlan` node that already uses it. - Expose a hit-aware cache API `LanceCache::get_or_insert_with_key_hit` returning `(Arc, was_cached)` and its `WeakLanceCache` counterpart, so per-query callers no longer have to wrap loaders in an extra `Arc`. - Aggregate into `ExecutionSummaryCounts` with an `index_cache_hit_ratio()` helper, propagate through Python `ScanStatistics` (pyi stubs updated) and Java `ScanStats` (10-arg constructor, 8-arg overload kept as `@Deprecated`). Instrumented boundaries in this release: BTree page, IVF partition (v2, write-cache scan path), inverted posting list (grouped and per-token) and its per-token metadata (`PostingMetadataKey`), inverted phrase positions (`PositionKey`), bitmap (Equals / Range / IsIn), ngram, and rtree page / null slot. Correctness fixes discovered during instrumentation: - Thread `Option<&dyn MetricsCollector>` through the whole BM25 stat path (`posting_len_for_token`, `df_for_term`, `bm25_scorer_for_final_tokens`, `bm25_stats_for_terms`, `bm25_base_scorer`, `build_global_bm25_scorer`) and forward it from `MatchQueryExec`, `FlatMatchFilterExec`, `FlatMatchQueryExec`, and the partition-local `bm25_search` fallback so per-token metadata cache lookups show up in per-query stats on every FTS path. - Bitmap: gate `load_bitmap` on `index_map.contains_key` before recording a miss, so a value not in the index short-circuits without polluting the counters, and thread `Some(metrics)` through the Range/IsIn code paths so their bitmap page reads finally count. - Inverted phrase positions: upgrade `read_positions` to `get_or_insert_with_key_hit` and record hit/miss under the caller's metrics context. - All cache-boundary sites record hit/miss even when the loader errors, matching the "loader ran" semantics documented on the accessors. Docs on `ExecutionSummaryCounts::index_cache_hits`, the `ScanStatistics` Python stub, and `ScanStats` Javadoc call out the instrumented boundary set, note that IVF streaming and legacy v1 IVF paths report as a miss on every call, and warn that a cold grouped posting-list lookup can record up to two misses for a single term. HNSW graph pages and quantizer codebooks are still uninstrumented and are noted as follow-up work. Tests: - `lance-index-core`: `local_metrics_collector_forwards_cache_counts`, `no_op_metrics_collector_ignores_cache_counts`. - `lance-index`: `test_page_cache_hit_miss_counts` in `btree.rs` runs a cold scan then a warm scan and asserts the expected 0/1 and 1/0 hit/miss counts; a targeted `get_or_insert_with_key_hit` test pins the `(value, was_cached)` contract. - Python: `test_btree_index_cache_hit_miss_stats` uses exact 1/0 and 0/1 counts; `test_bitmap_index_cache_hit_miss_stats` covers cold/warm IN queries plus a `color = 'purple'` absent-value regression that asserts hits/misses stay at zero; `test_phrase_query_cache_hit_miss_stats` exercises the `PositionKey` boundary through cold then warm phrase scans; `test_fts_index_cache_hit_miss_stats` proves the FTS wrapper forwarding. - Java: `ScannerTest.getStats()` reads `ScanStats.getIndexCacheHits()` / `getIndexCacheMisses()` to prove the 10-arg constructor wired up by the JNI patch does not throw and defaults to zero. --- java/lance-jni/src/blocking_scanner.rs | 4 +- .../main/java/org/lance/ipc/ScanStats.java | 84 +++++++ java/src/test/java/org/lance/ScannerTest.java | 4 + python/python/lance/lance/__init__.pyi | 25 ++ python/python/tests/test_scalar_index.py | 160 +++++++++++++ python/src/scanner.rs | 10 +- rust/lance-core/src/cache/mod.rs | 82 ++++++- rust/lance-datafusion/src/exec.rs | 83 ++++++- rust/lance-datafusion/src/utils.rs | 2 + rust/lance-index-core/src/metrics.rs | 117 ++++++++++ rust/lance-index/src/scalar/bitmap.rs | 114 ++++++++- rust/lance-index/src/scalar/btree.rs | 59 ++++- rust/lance-index/src/scalar/inverted.rs | 10 +- .../src/scalar/inverted/builder.rs | 5 +- rust/lance-index/src/scalar/inverted/index.rs | 217 ++++++++++++++---- rust/lance-index/src/scalar/ngram.rs | 9 +- rust/lance-index/src/scalar/rtree.rs | 24 +- rust/lance/src/index/vector/ivf/v2.rs | 13 +- rust/lance/src/io/exec/fts.rs | 33 ++- rust/lance/src/io/exec/utils.rs | 15 +- 20 files changed, 975 insertions(+), 95 deletions(-) diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 335cb2a4fa3..992776f60c8 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -637,7 +637,7 @@ fn inner_count_rows(env: &mut JNIEnv, j_scanner: JObject) -> Result { } const SCAN_STATS_CLASS: &str = "org/lance/ipc/ScanStats"; -const SCAN_STATS_CONSTRUCTOR_SIG: &str = "(JJJJJJLjava/util/Map;Ljava/util/Map;)V"; +const SCAN_STATS_CONSTRUCTOR_SIG: &str = "(JJJJJJJJLjava/util/Map;Ljava/util/Map;)V"; fn export_usize_map<'a>(env: &mut JNIEnv<'a>, map: &HashMap) -> Result> { let hash_map = env.new_object("java/util/HashMap", "()V", &[])?; @@ -669,6 +669,8 @@ impl IntoJava for &ExecutionSummaryCounts { JValueGen::Long(self.indices_loaded as i64), JValueGen::Long(self.parts_loaded as i64), JValueGen::Long(self.index_comparisons as i64), + JValueGen::Long(self.index_cache_hits() as i64), + JValueGen::Long(self.index_cache_misses() as i64), JValueGen::Object(&all_counts), JValueGen::Object(&all_times), ], diff --git a/java/src/main/java/org/lance/ipc/ScanStats.java b/java/src/main/java/org/lance/ipc/ScanStats.java index 24926d00911..e2160d0f371 100755 --- a/java/src/main/java/org/lance/ipc/ScanStats.java +++ b/java/src/main/java/org/lance/ipc/ScanStats.java @@ -30,6 +30,13 @@ public final class ScanStats { private final long indicesLoaded; private final long partsLoaded; private final long indexComparisons; + + /** Number of index cache page lookups served from memory in this scan. */ + private final long indexCacheHits; + + /** Number of index cache page lookups that had to load from storage in this scan. */ + private final long indexCacheMisses; + private final Map allCounts; private final Map allTimes; @@ -40,6 +47,8 @@ public ScanStats( long indicesLoaded, long partsLoaded, long indexComparisons, + long indexCacheHits, + long indexCacheMisses, Map allCounts, Map allTimes) { this.iops = iops; @@ -48,10 +57,42 @@ public ScanStats( this.indicesLoaded = indicesLoaded; this.partsLoaded = partsLoaded; this.indexComparisons = indexComparisons; + this.indexCacheHits = indexCacheHits; + this.indexCacheMisses = indexCacheMisses; this.allCounts = freezeMap(allCounts); this.allTimes = freezeMap(allTimes); } + /** + * Backwards-compatible constructor kept for existing callers that predate the addition of + * per-query index cache statistics. New code should use the 10-argument constructor that also + * accepts {@code indexCacheHits} and {@code indexCacheMisses}. + * + * @deprecated Use {@link #ScanStats(long, long, long, long, long, long, long, long, Map, Map)}. + */ + @Deprecated + public ScanStats( + long iops, + long requests, + long bytesRead, + long indicesLoaded, + long partsLoaded, + long indexComparisons, + Map allCounts, + Map allTimes) { + this( + iops, + requests, + bytesRead, + indicesLoaded, + partsLoaded, + indexComparisons, + 0L, + 0L, + allCounts, + allTimes); + } + private static Map freezeMap(Map map) { if (map == null || map.isEmpty()) { return Collections.emptyMap(); @@ -83,6 +124,41 @@ public long getIndexComparisons() { return indexComparisons; } + /** + * Number of index cache page lookups where the loader was not executed in this scan. + * + *

Counts both true cache hits on already-populated entries and coalesced concurrent loads (a + * follower attached to another caller's in-flight load). + * + *

Instrumented boundaries in this release: BTree, IVF v2 (write-cache scan path), inverted + * posting list (grouped and per-token) and its per-token metadata, inverted phrase positions, + * bitmap (Equals / Range / IsIn), ngram, rtree. + * + *

Caveats: + * + *

    + *
  • IVF v2 streaming scans and legacy v1 IVF partitions bypass the cache by design and are + * therefore reported as a miss on every call. + *
  • A cold posting-list lookup on the grouped inverted layout can record up to two misses + * (group + per-token metadata) for a single term. + *
+ * + *

Uninstrumented paths (HNSW graph pages, quantizer codebooks) do not contribute to either + * counter. See the sibling {@link #getIndexCacheMisses()} for the paired counter. + */ + public long getIndexCacheHits() { + return indexCacheHits; + } + + /** + * Number of index cache page lookups where the loader ran in this scan (the page was not resident + * and had to be materialised, typically from storage). See {@link #getIndexCacheHits()} for the + * paired counter and the list of instrumented boundaries. + */ + public long getIndexCacheMisses() { + return indexCacheMisses; + } + public Map getAllCounts() { return allCounts; } @@ -106,6 +182,8 @@ public boolean equals(Object o) { && indicesLoaded == that.indicesLoaded && partsLoaded == that.partsLoaded && indexComparisons == that.indexComparisons + && indexCacheHits == that.indexCacheHits + && indexCacheMisses == that.indexCacheMisses && Objects.equals(allCounts, that.allCounts) && Objects.equals(allTimes, that.allTimes); } @@ -119,6 +197,8 @@ public int hashCode() { indicesLoaded, partsLoaded, indexComparisons, + indexCacheHits, + indexCacheMisses, allCounts, allTimes); } @@ -138,6 +218,10 @@ public String toString() { + partsLoaded + ", indexComparisons=" + indexComparisons + + ", indexCacheHits=" + + indexCacheHits + + ", indexCacheMisses=" + + indexCacheMisses + ", allCounts=" + allCounts + ", allTimes=" diff --git a/java/src/test/java/org/lance/ScannerTest.java b/java/src/test/java/org/lance/ScannerTest.java index 0b07026bc68..3ca7481aa8c 100644 --- a/java/src/test/java/org/lance/ScannerTest.java +++ b/java/src/test/java/org/lance/ScannerTest.java @@ -212,6 +212,10 @@ void testDatasetScannerStats(@TempDir Path tempDir) throws Exception { assertTrue(statsOpt.isPresent()); ScanStats stats = statsOpt.get(); assertTrue(stats.getBytesRead() > 0 || !stats.getAllCounts().isEmpty()); + // Even without an index on this dataset, the two new counters must + // still marshal through JNI and default to zero rather than throwing. + assertTrue(stats.getIndexCacheHits() >= 0); + assertTrue(stats.getIndexCacheMisses() >= 0); } } } diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 9ba76cc592d..8c70fb48bc8 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -892,6 +892,31 @@ class ScanStatistics: indices_loaded: int parts_loaded: int index_comparisons: int + index_cache_hits: int + """Number of index cache page lookups where the loader was not executed + in this scan. Counts both true cache hits on already-populated entries + and coalesced concurrent loads (a follower attached to another caller's + in-flight load). + + Instrumented boundaries in this release: BTree, IVF v2 (write-cache scan + path), inverted posting list (grouped and per-token) and its per-token + metadata, inverted phrase positions, bitmap (Equals / Range / IsIn), + ngram, rtree. + + Caveats: + + * IVF v2 streaming scans and legacy v1 IVF partitions bypass the cache + by design and are therefore reported as a miss on every call. + * A cold posting-list lookup on the grouped inverted layout can record + up to two misses (group + per-token metadata) for a single term. + + Uninstrumented paths (HNSW graph pages, quantizer codebooks) do not + contribute to either counter.""" + index_cache_misses: int + """Number of index cache page lookups where the loader ran (the page was + not resident and had to be materialised, typically from storage). See + the sibling ``index_cache_hits`` for the paired counter and the list of + instrumented boundaries.""" all_counts: Dict[ str, int ] # Additional metrics for debugging purposes. Subject to change. diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index 43646a90722..bcb5cece18b 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -3346,6 +3346,166 @@ def scan_stats_callback(stats: lance.ScanStatistics): assert scan_stats.parts_loaded == 0 +def test_btree_index_cache_hit_miss_stats(tmp_path: Path): + """Cold scan reports index cache misses; warm scan reports hits. + + ScanStatistics.index_cache_{hits,misses} are populated at page-level cache + boundaries. On a freshly-loaded dataset the BTree page fetch must be a + miss; a second scan against the same in-memory Dataset re-uses the cached + page and therefore reports a hit with zero misses. + """ + scan_stats = None + + def scan_stats_callback(stats: lance.ScanStatistics): + nonlocal scan_stats + scan_stats = stats + + test_table = pa.table({"val": list(range(1000))}) + ds = lance.write_dataset(test_table, tmp_path) + ds.create_scalar_index("val", index_type="BTREE") + + # Reopen so the session cache starts cold. A single-key point lookup on + # a small dataset resolves to exactly one BTree page, so cold/warm counts + # are deterministic 1/0 and 0/1. + ds = lance.dataset(tmp_path) + ds.scanner(filter="val = 42", scan_stats_callback=scan_stats_callback).to_table() + assert scan_stats is not None + assert scan_stats.index_cache_misses == 1 + assert scan_stats.index_cache_hits == 0 + + # Same Dataset, warm cache — no new page loads, only hits. + ds.scanner(filter="val = 42", scan_stats_callback=scan_stats_callback).to_table() + assert scan_stats.index_cache_hits == 1 + assert scan_stats.index_cache_misses == 0 + + +def test_bitmap_index_cache_hit_miss_stats(tmp_path: Path): + """Bitmap Range/IN queries report cold misses and warm hits; a value + that is not in the index never reaches the loader and must not count. + + Guards against the ``BitmapIndex::search`` regressions where the + ``Range`` / ``IsIn`` branches used to drop the ``MetricsCollector`` (so + every lookup was silently ``0/0``), and where an equality on a value + absent from ``index_map`` recorded a spurious miss before short-circuiting + to the empty result. + """ + scan_stats = None + + def scan_stats_callback(stats: lance.ScanStatistics): + nonlocal scan_stats + scan_stats = stats + + test_table = pa.table({"color": ["red", "green", "blue", "yellow"] * 25}) + ds = lance.write_dataset(test_table, tmp_path) + ds.create_scalar_index("color", index_type="BITMAP") + + # Reopen so the session cache starts cold. + ds = lance.dataset(tmp_path) + ds.scanner( + filter="color IN ('red', 'blue')", scan_stats_callback=scan_stats_callback + ).to_table() + assert scan_stats is not None + assert scan_stats.index_cache_misses == 2 + assert scan_stats.index_cache_hits == 0 + + ds.scanner( + filter="color IN ('red', 'blue')", scan_stats_callback=scan_stats_callback + ).to_table() + assert scan_stats.index_cache_hits == 2 + assert scan_stats.index_cache_misses == 0 + + # A value that is not in the index short-circuits before the loader and + # must not touch either counter. + ds.scanner( + filter="color = 'purple'", scan_stats_callback=scan_stats_callback + ).to_table() + assert scan_stats.index_cache_hits == 0 + assert scan_stats.index_cache_misses == 0 + + +def test_phrase_query_cache_hit_miss_stats(tmp_path: Path): + """Phrase-query fallback populates ``PositionKey``; that boundary must + show up in per-query cache statistics. + + Guards against ``read_positions`` silently using the non-metric + ``get_or_insert_with_key`` API — before this fix, a warm phrase query + would report zero hits for the phrase-position cache slot even though + the loader was skipped. + + The cold path can already record a few hits (``bm25_stats_for_terms`` + populates ``PostingMetadataKey``, which is then re-read on the + posting-list path as a cross-boundary hit), so the cold assertion is + ``misses > hits`` rather than a strict zero. + """ + scan_stats = None + + def scan_stats_callback(stats: lance.ScanStatistics): + nonlocal scan_stats + scan_stats = stats + + test_table = pa.table( + {"text": ["quick brown fox jumps over lazy dog" for _ in range(50)]} + ) + ds = lance.write_dataset(test_table, tmp_path) + ds.create_scalar_index("text", index_type="INVERTED", with_position=True) + + ds = lance.dataset(tmp_path) + ds.scanner( + scan_stats_callback=scan_stats_callback, + full_text_query='"quick brown"', + ).to_table() + assert scan_stats is not None + assert scan_stats.index_cache_misses > 0 + assert scan_stats.index_cache_misses > scan_stats.index_cache_hits + + ds.scanner( + scan_stats_callback=scan_stats_callback, + full_text_query='"quick brown"', + ).to_table() + assert scan_stats.index_cache_hits > 0 + assert scan_stats.index_cache_misses == 0 + + +def test_fts_index_cache_hit_miss_stats(tmp_path: Path): + """Cold FTS scan reports misses; warm FTS scan reports hits. + + Guards the wrapper-forwarding fix in ``FtsIndexMetrics``: previously the + two new cache-hit/miss trait methods had default no-op implementations + that swallowed FTS-side events, so cache activity was reported as ``0/0`` + even for hot inverted-index scans. + + The cold path can still record a few hits when the same cache key is + read across boundaries in one query (e.g. ``bm25_stats_for_terms`` + populates ``PostingMetadataKey`` before ``posting_list`` re-reads it), + so the cold assertion is ``misses > hits`` rather than a strict zero. + """ + scan_stats = None + + def scan_stats_callback(stats: lance.ScanStatistics): + nonlocal scan_stats + scan_stats = stats + + test_table = pa.table({"fts": ["word" for _ in range(100)]}) + ds = lance.write_dataset(test_table, tmp_path) + ds.create_scalar_index("fts", index_type="INVERTED") + + # Reopen so the session cache starts cold. + ds = lance.dataset(tmp_path) + ds.scanner( + scan_stats_callback=scan_stats_callback, full_text_query="word" + ).to_table() + assert scan_stats is not None + assert scan_stats.index_cache_misses > 0 + assert scan_stats.index_cache_misses > scan_stats.index_cache_hits + + # Same Dataset, warm cache — posting-list / metadata reads must now hit. + ds.scanner( + scan_stats_callback=scan_stats_callback, full_text_query="word" + ).to_table() + assert scan_stats.index_cache_hits > 0 + assert scan_stats.index_cache_misses == 0 + + def test_fts_backward_v0_27_0(tmp_path: Path): path = ( Path(__file__).parent.parent.parent.parent diff --git a/python/src/scanner.rs b/python/src/scanner.rs index bbf1b3f35a3..8702537a340 100644 --- a/python/src/scanner.rs +++ b/python/src/scanner.rs @@ -66,6 +66,10 @@ pub struct ScanStatistics { pub parts_loaded: usize, /// Number of index comparisons performed pub index_comparisons: usize, + /// Number of index cache page lookups that were served from memory + pub index_cache_hits: usize, + /// Number of index cache page lookups that had to load from storage + pub index_cache_misses: usize, /// Additional metrics for more detailed statistics. These are subject to change in the future /// and should only be used for debugging purposes. pub all_counts: HashMap, @@ -80,6 +84,8 @@ impl ScanStatistics { indices_loaded: stats.indices_loaded, parts_loaded: stats.parts_loaded, index_comparisons: stats.index_comparisons, + index_cache_hits: stats.index_cache_hits(), + index_cache_misses: stats.index_cache_misses(), all_counts: stats.all_counts.clone(), } } @@ -89,13 +95,15 @@ impl ScanStatistics { impl ScanStatistics { fn __repr__(&self) -> String { format!( - "ScanStatistics(iops={}, requests={}, bytes_read={}, indices_loaded={}, parts_loaded={}, index_comparisons={}, all_counts={:?})", + "ScanStatistics(iops={}, requests={}, bytes_read={}, indices_loaded={}, parts_loaded={}, index_comparisons={}, index_cache_hits={}, index_cache_misses={}, all_counts={:?})", self.iops, self.requests, self.bytes_read, self.indices_loaded, self.parts_loaded, self.index_comparisons, + self.index_cache_hits, + self.index_cache_misses, self.all_counts ) } diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index dc64b79f047..9ff89e6d34f 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -370,6 +370,37 @@ impl LanceCache { cache_key: K, loader: F, ) -> Result> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + F: FnOnce() -> Fut + Send, + Fut: Future> + Send, + { + self.get_or_insert_with_key_hit(cache_key, loader) + .await + .map(|(v, _)| v) + } + + /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but + /// also returns a boolean indicating whether the loader was skipped for + /// this call. + /// + /// - `true` means this call did **not** execute the loader. That covers + /// both a true cache hit on an already-populated entry and a coalesced + /// concurrent load where an in-flight loader started by a different + /// caller produced the value. + /// - `false` means the loader ran on this call (a real cache miss). + /// + /// Callers that want strict "served from cache" semantics should treat + /// coalesced loads as misses; the current backend does not distinguish the + /// two cases. Prefer this over rolling a caller-side `Arc` + /// when the caller needs per-query hit/miss counters — the backend already + /// tracks this bit internally and this method just exposes it. + pub async fn get_or_insert_with_key_hit( + &self, + cache_key: K, + loader: F, + ) -> Result<(Arc, bool)> where K: CacheKey, K::ValueType: DeepSizeOf + Send + Sync + 'static, @@ -396,7 +427,7 @@ impl LanceCache { self.misses.fetch_add(1, Ordering::Relaxed); } - Ok(entry.downcast::().unwrap()) + Ok((entry.downcast::().unwrap(), was_cached)) } pub async fn insert_unsized_with_key(&self, cache_key: &K, metadata: Arc) @@ -500,6 +531,26 @@ impl WeakLanceCache { cache_key: K, loader: F, ) -> Result> + where + K: CacheKey, + K::ValueType: DeepSizeOf + Send + Sync + 'static, + F: FnOnce() -> Fut + Send, + Fut: Future> + Send, + { + self.get_or_insert_with_key_hit(cache_key, loader) + .await + .map(|(v, _)| v) + } + + /// Same as [`get_or_insert_with_key`](Self::get_or_insert_with_key), but + /// also returns a boolean indicating whether the loader was skipped for + /// this call. See [`LanceCache::get_or_insert_with_key_hit`] for the + /// coalesced-load caveat. + pub async fn get_or_insert_with_key_hit( + &self, + cache_key: K, + loader: F, + ) -> Result<(Arc, bool)> where K: CacheKey, K::ValueType: DeepSizeOf + Send + Sync + 'static, @@ -520,10 +571,10 @@ impl WeakLanceCache { } else { self.misses.fetch_add(1, Ordering::Relaxed); } - Ok(entry.downcast::().unwrap()) + Ok((entry.downcast::().unwrap(), was_cached)) } else { log::warn!("WeakLanceCache: cache no longer available, computing without caching"); - loader().await.map(Arc::new) + loader().await.map(|v| (Arc::new(v), false)) } } @@ -954,6 +1005,31 @@ mod tests { ); } + #[tokio::test] + async fn test_cache_get_or_insert_with_key_hit() { + let cache = LanceCache::with_capacity(1000); + + // Cold: loader runs, was_cached = false. + let (v, was_cached) = cache + .get_or_insert_with_key_hit(TestKey::>::new("k"), || async { + Ok(vec![1, 2, 3]) + }) + .await + .unwrap(); + assert_eq!(*v, vec![1, 2, 3]); + assert!(!was_cached); + + // Warm: loader must not run and was_cached = true. + let (v, was_cached) = cache + .get_or_insert_with_key_hit(TestKey::>::new("k"), || async { + panic!("should not be called") + }) + .await + .unwrap(); + assert_eq!(*v, vec![1, 2, 3]); + assert!(was_cached); + } + #[tokio::test] async fn test_custom_backend() { use async_trait::async_trait; diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index c9822250651..994e16b783d 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -58,8 +58,9 @@ use crate::udf::register_functions; use crate::{ chunker::StrictBatchSizeStream, utils::{ - BYTES_READ_METRIC, INDEX_COMPARISONS_METRIC, INDICES_LOADED_METRIC, IOPS_METRIC, - MetricsExt, PARTS_LOADED_METRIC, REQUESTS_METRIC, + BYTES_READ_METRIC, INDEX_CACHE_HITS_METRIC, INDEX_CACHE_MISSES_METRIC, + INDEX_COMPARISONS_METRIC, INDICES_LOADED_METRIC, IOPS_METRIC, MetricsExt, + PARTS_LOADED_METRIC, REQUESTS_METRIC, }, }; @@ -490,12 +491,88 @@ pub struct ExecutionSummaryCounts { pub index_comparisons: usize, /// Additional metrics for more detailed statistics. These are subject to change in the future /// and should only be used for debugging purposes. + /// + /// Newer metrics (e.g. [`INDEX_CACHE_HITS_METRIC`], [`INDEX_CACHE_MISSES_METRIC`]) are added + /// here rather than as `pub` fields, so this struct stays backwards compatible for callers + /// that construct or destructure it. Prefer the typed accessors below. pub all_counts: HashMap, /// Additional time metrics for more detailed statistics, stored in nanoseconds. /// These are subject to change in the future and should only be used for debugging purposes. pub all_times: HashMap, } +impl ExecutionSummaryCounts { + /// Number of index cache page lookups where the loader was not executed + /// (per-page granularity). + /// + /// A "hit" is any page-level lookup at an instrumented cache boundary that + /// did not run the loader on this call. That covers both a true cache hit + /// on an already-populated entry and a coalesced concurrent load where an + /// in-flight loader started by a different caller produced the value. + /// + /// Instrumented boundaries in this release: + /// BTree page, IVF partition (v2, `write_cache=true` scan path), inverted + /// posting list (grouped and per-token), inverted per-token metadata + /// (`PostingMetadataKey`), inverted phrase positions (`PositionKey`), + /// bitmap posting (Equals / Range / IsIn), ngram posting, and rtree page + /// / null slot. + /// + /// Caveats: + /// * IVF v2 streaming scans and legacy v1 IVF partitions run + /// `load_partition` with `write_cache=false`. Those loads always execute + /// the loader and never write the result back, so they are reported as a + /// miss on every call. See [`Self::index_cache_hit_ratio`]. + /// * A cold posting-list lookup on the grouped inverted layout can record + /// up to two misses (posting-list group + per-token metadata) for a + /// single term. + /// + /// Other index cache boundaries such as HNSW graph pages and quantizer + /// codebooks are not yet instrumented; a scan that only touches those + /// paths returns `0` here. + pub fn index_cache_hits(&self) -> usize { + self.all_counts + .get(INDEX_CACHE_HITS_METRIC) + .copied() + .unwrap_or(0) + } + + /// Number of index cache page lookups that had to execute the loader + /// (per-page granularity). + /// + /// A "miss" is any page-level lookup at an instrumented cache boundary + /// where the loader ran, i.e. the page was not resident and had to be + /// materialised (typically from storage). See + /// [`Self::index_cache_hits`] for the paired counter and the list of + /// instrumented boundaries. + pub fn index_cache_misses(&self) -> usize { + self.all_counts + .get(INDEX_CACHE_MISSES_METRIC) + .copied() + .unwrap_or(0) + } + + /// Ratio of index cache hits to total lookups. Returns `0.0` when no lookups + /// were recorded in this scan. + /// + /// This ratio only reflects paths that write their result back to the + /// index cache. Streaming scans (IVF v2 `write_cache=false` and legacy v1 + /// IVF `load_partition_stream`) intentionally bypass the cache and are + /// counted as misses on every call, so a workload dominated by streaming + /// vector scans will report a hit ratio near `0.0` regardless of cache + /// size. + pub fn index_cache_hit_ratio(&self) -> f32 { + // Widen to u128 before summing so a pathological (hits + misses) + // overflow can't panic in debug builds nor wrap in release builds. + let hits = self.index_cache_hits() as u128; + let total = hits + self.index_cache_misses() as u128; + if total == 0 { + 0.0 + } else { + hits as f32 / total as f32 + } + } +} + pub fn collect_execution_metrics(node: &dyn ExecutionPlan, counts: &mut ExecutionSummaryCounts) { if let Some(metrics) = node.metrics() { for (metric_name, count) in metrics.iter_counts() { @@ -556,6 +633,8 @@ fn report_plan_summary_metrics(plan: &dyn ExecutionPlan, options: &LanceExecutio indices_loaded = counts.indices_loaded, parts_loaded = counts.parts_loaded, index_comparisons = counts.index_comparisons, + index_cache_hits = counts.index_cache_hits(), + index_cache_misses = counts.index_cache_misses(), ); } if let Some(callback) = options.execution_stats_callback.as_ref() { diff --git a/rust/lance-datafusion/src/utils.rs b/rust/lance-datafusion/src/utils.rs index 470831bf0bc..e660c8ee44d 100644 --- a/rust/lance-datafusion/src/utils.rs +++ b/rust/lance-datafusion/src/utils.rs @@ -244,6 +244,8 @@ pub const INDICES_LOADED_METRIC: &str = "indices_loaded"; pub const PARTS_LOADED_METRIC: &str = "parts_loaded"; pub const PARTITIONS_RANKED_METRIC: &str = "partitions_ranked"; pub const INDEX_COMPARISONS_METRIC: &str = "index_comparisons"; +pub const INDEX_CACHE_HITS_METRIC: &str = "index_cache_hits"; +pub const INDEX_CACHE_MISSES_METRIC: &str = "index_cache_misses"; pub const FRAGMENTS_SCANNED_METRIC: &str = "fragments_scanned"; pub const RANGES_SCANNED_METRIC: &str = "ranges_scanned"; pub const ROWS_SCANNED_METRIC: &str = "rows_scanned"; diff --git a/rust/lance-index-core/src/metrics.rs b/rust/lance-index-core/src/metrics.rs index 8c0c119a3c3..aaa016a5dda 100644 --- a/rust/lance-index-core/src/metrics.rs +++ b/rust/lance-index-core/src/metrics.rs @@ -49,6 +49,28 @@ pub trait MetricsCollector: Send + Sync { /// The goal is to provide some visibility into the compute cost of the search fn record_comparisons(&self, num_comparisons: usize); + /// Record index cache hits observed while serving this query. + /// + /// A "hit" is one page-level lookup (partition, posting list, BTree page, etc.) + /// that was served from the in-memory index cache without touching storage. + fn record_index_cache_hits(&self, _num_hits: usize) {} + + /// Convenience for a single cache hit. + fn record_index_cache_hit(&self) { + self.record_index_cache_hits(1); + } + + /// Record index cache misses observed while serving this query. + /// + /// A "miss" is one page-level lookup that had to be loaded from storage + /// because it was not present in the cache. + fn record_index_cache_misses(&self, _num_misses: usize) {} + + /// Convenience for a single cache miss. + fn record_index_cache_miss(&self) { + self.record_index_cache_misses(1); + } + /// Record AND candidates returned from WAND alignment to the scoring loop. /// /// This excludes candidates pruned before `next()` returns. Use this with @@ -91,6 +113,12 @@ pub struct LocalMetricsCollector { pub parts_loaded: AtomicUsize, pub index_loads: AtomicUsize, pub comparisons: AtomicUsize, + // Kept `pub(crate)` so that adding new metric fields to this public struct + // does not break downstream callers that construct or destructure the + // existing three fields. Callers can still read cumulative values via + // [`Self::index_cache_hits`] / [`Self::index_cache_misses`]. + pub(crate) index_cache_hits: AtomicUsize, + pub(crate) index_cache_misses: AtomicUsize, } impl LocalMetricsCollector { @@ -98,6 +126,18 @@ impl LocalMetricsCollector { other.record_parts_loaded(self.parts_loaded.load(Ordering::Relaxed)); other.record_index_loads(self.index_loads.load(Ordering::Relaxed)); other.record_comparisons(self.comparisons.load(Ordering::Relaxed)); + other.record_index_cache_hits(self.index_cache_hits.load(Ordering::Relaxed)); + other.record_index_cache_misses(self.index_cache_misses.load(Ordering::Relaxed)); + } + + /// Cumulative index cache hits recorded so far. + pub fn index_cache_hits(&self) -> usize { + self.index_cache_hits.load(Ordering::Relaxed) + } + + /// Cumulative index cache misses recorded so far. + pub fn index_cache_misses(&self) -> usize { + self.index_cache_misses.load(Ordering::Relaxed) } } @@ -114,4 +154,81 @@ impl MetricsCollector for LocalMetricsCollector { self.comparisons .fetch_add(num_comparisons, Ordering::Relaxed); } + + fn record_index_cache_hits(&self, num_hits: usize) { + self.index_cache_hits.fetch_add(num_hits, Ordering::Relaxed); + } + + fn record_index_cache_misses(&self, num_misses: usize) { + self.index_cache_misses + .fetch_add(num_misses, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct SumSink { + parts: AtomicUsize, + loads: AtomicUsize, + comparisons: AtomicUsize, + hits: AtomicUsize, + misses: AtomicUsize, + } + + impl MetricsCollector for SumSink { + fn record_parts_loaded(&self, n: usize) { + self.parts.fetch_add(n, Ordering::Relaxed); + } + fn record_index_loads(&self, n: usize) { + self.loads.fetch_add(n, Ordering::Relaxed); + } + fn record_comparisons(&self, n: usize) { + self.comparisons.fetch_add(n, Ordering::Relaxed); + } + fn record_index_cache_hits(&self, n: usize) { + self.hits.fetch_add(n, Ordering::Relaxed); + } + fn record_index_cache_misses(&self, n: usize) { + self.misses.fetch_add(n, Ordering::Relaxed); + } + } + + #[test] + fn local_metrics_collector_forwards_cache_counts() { + let local = LocalMetricsCollector::default(); + local.record_index_cache_hit(); + local.record_index_cache_hit(); + local.record_index_cache_misses(3); + local.record_part_load(); + local.record_index_load(); + local.record_comparisons(5); + + let sink = SumSink { + parts: AtomicUsize::new(0), + loads: AtomicUsize::new(0), + comparisons: AtomicUsize::new(0), + hits: AtomicUsize::new(0), + misses: AtomicUsize::new(0), + }; + local.dump_into(&sink); + + assert_eq!(sink.parts.load(Ordering::Relaxed), 1); + assert_eq!(sink.loads.load(Ordering::Relaxed), 1); + assert_eq!(sink.comparisons.load(Ordering::Relaxed), 5); + assert_eq!(sink.hits.load(Ordering::Relaxed), 2); + assert_eq!(sink.misses.load(Ordering::Relaxed), 3); + } + + #[test] + fn no_op_metrics_collector_ignores_cache_counts() { + // Ensures existing implementors that do not override cache-count methods + // remain sound (default impl is a no-op). + let collector = NoOpMetricsCollector; + collector.record_index_cache_hit(); + collector.record_index_cache_miss(); + collector.record_index_cache_hits(10); + collector.record_index_cache_misses(20); + } } diff --git a/rust/lance-index/src/scalar/bitmap.rs b/rust/lance-index/src/scalar/bitmap.rs index 91595b277f0..04bafb3800b 100644 --- a/rust/lance-index/src/scalar/bitmap.rs +++ b/rust/lance-index/src/scalar/bitmap.rs @@ -446,22 +446,31 @@ impl BitmapIndex { return Ok(self.null_map.clone()); } + // A value that isn't in `index_map` never reaches the loader or the + // cache, so it should not touch the per-query cache counters either. + // Checking here (before the cached-lookup fast path) also avoids + // returning an unmapped-value response as a spurious cache hit if a + // prior insert somehow ended up under `cache_key`. + let row_offset = match self.index_map.get(key) { + Some(loc) => *loc, + None => return Ok(Arc::new(RowAddrTreeMap::default())), + }; + let cache_key = BitmapKey { value: key.clone() }; if let Some(cached) = self.index_cache.get_with_key(&cache_key).await { + if let Some(metrics) = metrics { + metrics.record_index_cache_hit(); + } return Ok(cached); } // Record that we're loading a partition from disk if let Some(metrics) = metrics { + metrics.record_index_cache_miss(); metrics.record_part_load(); } - let row_offset = match self.index_map.get(key) { - Some(loc) => *loc, - None => return Ok(Arc::new(RowAddrTreeMap::default())), - }; - let page_lookup_file = self.lazy_reader.get().await?; let batch = page_lookup_file .read_range(row_offset..row_offset + 1, Some(&["bitmaps"])) @@ -698,7 +707,7 @@ impl ScalarIndex for BitmapIndex { } else { let bitmaps: Vec<_> = stream::iter( keys.into_iter() - .map(|key| async move { self.load_bitmap(&key, None).await }), + .map(|key| async move { self.load_bitmap(&key, Some(metrics)).await }), ) .buffer_unordered(get_num_compute_intensive_cpus()) .try_collect() @@ -740,7 +749,7 @@ impl ScalarIndex for BitmapIndex { // Load bitmaps in parallel let mut bitmaps: Vec<_> = stream::iter( keys.into_iter() - .map(|key| async move { self.load_bitmap(&key, None).await }), + .map(|key| async move { self.load_bitmap(&key, Some(metrics)).await }), ) .buffer_unordered(get_num_compute_intensive_cpus()) .try_collect() @@ -1878,7 +1887,7 @@ impl ScalarIndexPlugin for BitmapIndexPlugin { #[cfg(test)] mod tests { use super::*; - use crate::metrics::NoOpMetricsCollector; + use crate::metrics::{LocalMetricsCollector, NoOpMetricsCollector}; use crate::scalar::lance_format::LanceIndexStore; use arrow_array::{RecordBatch, StringArray, UInt64Array, record_batch}; use arrow_schema::{DataType, Field, Schema}; @@ -2128,6 +2137,95 @@ mod tests { } } + /// Regression test for the review fix that gates `load_bitmap` on + /// `index_map.contains_key` before recording a miss: a value that is + /// not present in the index must short-circuit before touching the + /// per-query cache counters. Previously an Equals query for a missing + /// value would silently bump `index_cache_misses` and `parts_loaded` + /// on every call even though no bitmap page was actually loaded. + #[tokio::test] + async fn test_bitmap_absent_value_records_no_cache_activity() { + let tmpdir = TempObjDir::default(); + let store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let colors = vec!["red", "blue", "green", "yellow"]; + let row_ids = (0u64..4u64).collect::>(); + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Utf8, false), + Field::new("_rowid", DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(colors)), + Arc::new(UInt64Array::from(row_ids)), + ], + ) + .unwrap(); + let batch = sort_batch_by_value(&batch); + let stream = stream::once(async move { Ok(batch) }); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + BitmapIndexPlugin::train_bitmap_index(stream, store.as_ref()) + .await + .unwrap(); + + // Keep the `LanceCache` alive in test scope so the `WeakLanceCache` + // inside `BitmapIndex` can upgrade during search. + let cache = LanceCache::with_capacity(1024 * 1024); + let index = BitmapIndex::load(store.clone(), None, &cache) + .await + .unwrap(); + + // Equals on a value that is not in `index_map` must not touch + // the cache counters and must not report a part load. + let metrics = LocalMetricsCollector::default(); + let query = SargableQuery::Equals(ScalarValue::Utf8(Some("purple".to_string()))); + let result = index.search(&query, &metrics).await.unwrap(); + if let SearchResult::Exact(row_ids) = result { + assert!(row_ids.true_rows().is_empty()); + } else { + panic!("Expected exact search result"); + } + assert_eq!( + metrics.index_cache_hits(), + 0, + "absent value must not record any cache hits", + ); + assert_eq!( + metrics.index_cache_misses(), + 0, + "absent value must not record a cache miss (no loader ran)", + ); + + // IsIn covering only absent values also stays at 0/0. + let metrics = LocalMetricsCollector::default(); + let query = SargableQuery::IsIn(vec![ + ScalarValue::Utf8(Some("purple".to_string())), + ScalarValue::Utf8(Some("teal".to_string())), + ]); + let result = index.search(&query, &metrics).await.unwrap(); + if let SearchResult::Exact(row_ids) = result { + assert!(row_ids.true_rows().is_empty()); + } else { + panic!("Expected exact search result"); + } + assert_eq!(metrics.index_cache_hits(), 0); + assert_eq!(metrics.index_cache_misses(), 0); + + // Sanity: a present value on the same cold cache still records + // exactly one miss, proving the counters are wired up and the + // absent-value path above is not silently no-op. + let metrics = LocalMetricsCollector::default(); + let query = SargableQuery::Equals(ScalarValue::Utf8(Some("red".to_string()))); + index.search(&query, &metrics).await.unwrap(); + assert_eq!(metrics.index_cache_hits(), 0); + assert_eq!(metrics.index_cache_misses(), 1); + } + // Regression test for the O(N log N) warm-cache rebuild introduced in // commit 4de5ce67d. BitmapIndexState now caches the parsed Arc // so that get_from_cache skips parse_lookup_batch on warm hits. diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 5bf89dfa29f..186a3813b73 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1635,11 +1635,17 @@ impl BTreeIndex { index_reader: LazyIndexReader, metrics: &dyn MetricsCollector, ) -> Result> { - self.index_cache - .get_or_insert_with_key(BTreePageKey { page_number }, move || async move { + let result = self + .index_cache + .get_or_insert_with_key_hit(BTreePageKey { page_number }, move || async move { self.read_page(page_number, index_reader, metrics).await }) - .await + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + result.map(|(page, _)| page) } #[instrument(level = "debug", skip_all)] @@ -3704,6 +3710,53 @@ mod tests { assert_eq!(metrics.parts_loaded.load(Ordering::Relaxed), 1); } + #[tokio::test] + async fn test_page_cache_hit_miss_counts() { + let tmpdir = TempObjDir::default(); + let test_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + let data = gen_batch() + .col("value", array::step::()) + .col("_rowid", array::step::()) + .into_df_exec(RowCount::from(1000), BatchCount::from(10)); + let schema = data.schema(); + let sort_expr = PhysicalSortExpr::new_default(col("value", schema.as_ref()).unwrap()); + let plan = Arc::new(SortExec::new([sort_expr].into(), data)); + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let stream = break_stream(stream, 64); + let stream = stream.map_err(DataFusionError::from); + let stream = + Box::pin(RecordBatchStreamAdapter::new(schema, stream)) as SendableRecordBatchStream; + + train_btree_index(stream, test_store.as_ref(), 64, None, None) + .await + .unwrap(); + + let cache = Arc::new(LanceCache::with_capacity(100 * 1024 * 1024)); + let index = BTreeIndex::load(test_store, None, cache.as_ref()) + .await + .unwrap(); + + // First search: cold cache — the page fetch must miss. + let query = SargableQuery::Equals(ScalarValue::Float32(Some(0.0))); + let cold = LocalMetricsCollector::default(); + index.search(&query, &cold).await.unwrap(); + assert_eq!(cold.index_cache_hits(), 0); + assert_eq!(cold.index_cache_misses(), 1); + assert_eq!(cold.parts_loaded.load(Ordering::Relaxed), 1); + + // Second search: same key, page must now be served from cache. + let warm = LocalMetricsCollector::default(); + index.search(&query, &warm).await.unwrap(); + assert_eq!(warm.index_cache_hits(), 1); + assert_eq!(warm.index_cache_misses(), 0); + assert_eq!(warm.parts_loaded.load(Ordering::Relaxed), 0); + } + #[tokio::test] async fn test_like_prefix_search() { use arrow::datatypes::DataType; diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 72604243c1c..f51ff18a103 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -74,6 +74,11 @@ fn scorer_terms( /// statistics rather than per-segment statistics. Computes the union of /// fuzzy-expanded terms when `params.fuzziness` is set. /// +/// `metrics`, when provided, is forwarded to the per-token metadata cache +/// boundary on each segment so callers running under an `ExecutionPlan` +/// (e.g. `MatchQueryExec`) see the reads triggered here in their per-query +/// `index_cache_hits`/`index_cache_misses` counters. +/// /// Public as the canonical producer paired with the `with_base_scorer` /// consumer on FTS exec types: callers holding `Arc` segment /// handles locally can construct an injectable scorer without reimplementing @@ -84,13 +89,14 @@ pub async fn build_global_bm25_scorer( indices: &[Arc], query_tokens: &Tokens, params: &FtsSearchParams, + metrics: Option<&dyn crate::scalar::MetricsCollector>, ) -> Result { let terms = scorer_terms(indices, query_tokens, params)?; let first_index = indices.first().ok_or_else(|| { lance_core::Error::invalid_input("FTS index requires at least one segment") })?; let (mut total_tokens, mut num_docs, first_token_docs) = - first_index.bm25_stats_for_terms(&terms).await?; + first_index.bm25_stats_for_terms(&terms, metrics).await?; let mut token_docs = HashMap::with_capacity(terms.len()); for (term, count) in terms.iter().cloned().zip(first_token_docs) { token_docs.insert(term, count); @@ -98,7 +104,7 @@ pub async fn build_global_bm25_scorer( for index in indices.iter().skip(1) { let (segment_total_tokens, segment_num_docs, segment_token_docs) = - index.bm25_stats_for_terms(&terms).await?; + index.bm25_stats_for_terms(&terms, metrics).await?; total_tokens += segment_total_tokens; num_docs += segment_num_docs; for (term, count) in terms.iter().zip(segment_token_docs) { diff --git a/rust/lance-index/src/scalar/inverted/builder.rs b/rust/lance-index/src/scalar/inverted/builder.rs index c90f96053ef..588c8822942 100644 --- a/rust/lance-index/src/scalar/inverted/builder.rs +++ b/rust/lance-index/src/scalar/inverted/builder.rs @@ -3565,8 +3565,9 @@ mod tests { .await?; let index = InvertedIndex::load(store, None, &LanceCache::no_cache()).await?; - let (total_tokens, num_docs, token_docs) = - index.bm25_stats_for_terms(&["hello".to_string()]).await?; + let (total_tokens, num_docs, token_docs) = index + .bm25_stats_for_terms(&["hello".to_string()], None) + .await?; assert_eq!(total_tokens, 1); assert_eq!(num_docs, 1); assert_eq!(token_docs, vec![1]); diff --git a/rust/lance-index/src/scalar/inverted/index.rs b/rust/lance-index/src/scalar/inverted/index.rs index 368660133ca..a6db0e355c9 100644 --- a/rust/lance-index/src/scalar/inverted/index.rs +++ b/rust/lance-index/src/scalar/inverted/index.rs @@ -656,16 +656,21 @@ impl InvertedIndex { /// expansions, not just the raw query tokens — otherwise /// `query_weight(expanded_token)` returns 0 and the BM25 contribution /// of every expanded match is discarded. + /// + /// `metrics` is forwarded to the per-token metadata cache boundary so the + /// caller's per-query cache counters see the reads triggered here. pub async fn bm25_base_scorer( &self, query_tokens: &Tokens, params: &FtsSearchParams, + metrics: Option<&dyn MetricsCollector>, ) -> Result { if matches!(params.fuzziness, Some(n) if n != 0) { let expanded = self.expand_fuzzy_tokens(query_tokens, params)?; - self.bm25_scorer_for_final_tokens(&expanded).await + self.bm25_scorer_for_final_tokens(&expanded, metrics).await } else { - self.bm25_scorer_for_final_tokens(query_tokens).await + self.bm25_scorer_for_final_tokens(query_tokens, metrics) + .await } } @@ -673,7 +678,11 @@ impl InvertedIndex { /// the terms and pull their document frequencies. `bm25_search` calls /// this with the tokens it already expanded, so the expansion runs once /// per query rather than once for the scorer and once per partition. - async fn bm25_scorer_for_final_tokens(&self, tokens: &Tokens) -> Result { + async fn bm25_scorer_for_final_tokens( + &self, + tokens: &Tokens, + metrics: Option<&dyn MetricsCollector>, + ) -> Result { let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; let mut terms: Vec = Vec::new(); let mut seen = HashSet::new(); @@ -684,16 +693,25 @@ impl InvertedIndex { } let mut token_docs = HashMap::with_capacity(terms.len()); for term in &terms { - let df = self.df_for_term(term).await?; + let df = self.df_for_term(term, metrics).await?; token_docs.insert(term.clone(), df); } Ok(MemBM25Scorer::new(total_tokens, num_docs, token_docs)) } - pub async fn bm25_stats_for_terms(&self, terms: &[String]) -> Result<(u64, usize, Vec)> { + /// Collect the `(total_tokens, num_docs, per_term_df)` triple used to + /// combine per-segment BM25 statistics into a global scorer. `metrics` + /// is threaded to record per-token metadata cache activity in the + /// caller's per-query counters. + pub async fn bm25_stats_for_terms( + &self, + terms: &[String], + metrics: Option<&dyn MetricsCollector>, + ) -> Result<(u64, usize, Vec)> { let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?; let token_docs = - futures::future::try_join_all(terms.iter().map(|term| self.df_for_term(term))).await?; + futures::future::try_join_all(terms.iter().map(|term| self.df_for_term(term, metrics))) + .await?; Ok((total_tokens, num_docs, token_docs)) } @@ -728,7 +746,11 @@ impl InvertedIndex { /// Sum the posting-list length for `term` across this index's partitions /// via single-row reads, with partition lookups bounded by the store's /// `io_parallelism()`. - async fn df_for_term(&self, term: &str) -> Result { + async fn df_for_term( + &self, + term: &str, + metrics: Option<&dyn MetricsCollector>, + ) -> Result { let io_parallelism = self.store.io_parallelism(); let futures = self .partitions @@ -737,7 +759,11 @@ impl InvertedIndex { let part = part.clone(); async move { match part.tokens.get(term) { - Some(token_id) => part.inverted_list.posting_len_for_token(token_id).await, + Some(token_id) => { + part.inverted_list + .posting_len_for_token(token_id, metrics) + .await + } None => Ok(0), } } @@ -846,7 +872,9 @@ impl InvertedIndex { let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer { base_scorer } else { - local_scorer = self.bm25_scorer_for_final_tokens(tokens.as_ref()).await?; + local_scorer = self + .bm25_scorer_for_final_tokens(tokens.as_ref(), Some(metrics.as_ref())) + .await?; &local_scorer }; let impact_scorer = Arc::new(scorer.clone()); @@ -3089,14 +3117,22 @@ impl PostingListReader { /// not been loaded yet, and never triggers the bulk load itself. The stats /// path uses this so a single-term `df` lookup costs O(1) bytes rather /// than O(num_unique_tokens). - pub(crate) async fn posting_len_for_token(&self, token_id: u32) -> Result { + /// + /// `metrics` is threaded through so callers holding a real + /// `MetricsCollector` (e.g. `MatchQueryExec`) record the per-token + /// `PostingMetadataKey` cache boundary in their per-query counters. + pub(crate) async fn posting_len_for_token( + &self, + token_id: u32, + metrics: Option<&dyn MetricsCollector>, + ) -> Result { match &self.metadata { PostingMetadata::LegacyV1 { .. } => Ok(self.posting_len(token_id)), PostingMetadata::V2 { metadata } => { if let Some(metadata) = metadata.get() { return Ok(metadata.lengths[token_id as usize] as usize); } - let (_, length) = self.posting_metadata_for_token(token_id).await?; + let (_, length) = self.posting_metadata_for_token(token_id, metrics).await?; length .map(|len| len as usize) .ok_or_else(|| Error::index("posting length metadata missing".to_string())) @@ -3113,6 +3149,7 @@ impl PostingListReader { pub(crate) async fn posting_metadata_for_token( &self, token_id: u32, + metrics: Option<&dyn MetricsCollector>, ) -> Result<(Option, Option)> { match &self.metadata { PostingMetadata::LegacyV1 { max_scores, .. } => { @@ -3125,9 +3162,9 @@ impl PostingListReader { Some(loaded.lengths[token_id as usize]), )); } - let metadata = self + let result = self .index_cache - .get_or_insert_with_key(PostingMetadataKey { token_id }, || async move { + .get_or_insert_with_key_hit(PostingMetadataKey { token_id }, || async move { let token_id = token_id as usize; let batch = self .reader @@ -3137,7 +3174,14 @@ impl PostingListReader { let length = batch[LENGTH_COL].as_primitive::().value(0); Ok(PostingMetadataValue { max_score, length }) }) - .await?; + .await; + if let Some(metrics) = metrics { + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + } + let metadata = result.map(|(v, _)| v)?; Ok((Some(metadata.max_score), Some(metadata.length))) } } @@ -3245,9 +3289,9 @@ impl PostingListReader { // Grouped path (issue #7040): one cache entry covers rows // [start, end), so neighbouring rare terms share a single read. Some((start, end)) => { - let group = self + let result = self .index_cache - .get_or_insert_with_key( + .get_or_insert_with_key_hit( posting_list_group_cache_key(start, end, self.has_impacts), || async move { metrics.record_part_load(); @@ -3255,9 +3299,15 @@ impl PostingListReader { self.load_posting_list_group(start, end).await }, ) - .await?; + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + let (group, _) = result?; let (max_score, length) = if group.needs_external_metadata() { - self.posting_metadata_for_token(token_id).await? + self.posting_metadata_for_token(token_id, Some(metrics)) + .await? } else { (None, None) }; @@ -3272,32 +3322,37 @@ impl PostingListReader { } // Fallback for layouts that cannot use row-based groups: one cache // entry per token. - None => self - .index_cache - .get_or_insert_with_key( - posting_list_cache_key(token_id, self.has_impacts), - || async move { - metrics.record_part_load(); - info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); - // Fetch the posting batch and this token's (max_score, - // length) in parallel; for cold v2 partitions this is one - // single-row metadata read plus one posting-row read, - // instead of pulling the full per-token metadata table. - let (batch, (max_score, length)) = futures::try_join!( - self.posting_batch(token_id, false), - self.posting_metadata_for_token(token_id), - )?; - self.posting_list_from_batch(&batch, max_score, length) - }, - ) - .await? - .as_ref() - .clone(), + None => { + let result = self + .index_cache + .get_or_insert_with_key_hit( + posting_list_cache_key(token_id, self.has_impacts), + || async move { + metrics.record_part_load(); + info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="inverted", part_id=token_id); + // Fetch the posting batch and this token's (max_score, + // length) in parallel; for cold v2 partitions this is one + // single-row metadata read plus one posting-row read, + // instead of pulling the full per-token metadata table. + let (batch, (max_score, length)) = futures::try_join!( + self.posting_batch(token_id, false), + self.posting_metadata_for_token(token_id, Some(metrics)), + )?; + self.posting_list_from_batch(&batch, max_score, length) + }, + ) + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + result?.0.as_ref().clone() + } }; if is_phrase_query && !posting.has_position() { // hit the cache and when the cache was populated, the positions column was not loaded - let positions = self.read_positions(token_id).await?; + let positions = self.read_positions(token_id, metrics).await?; posting.set_positions(positions); } @@ -3800,8 +3855,12 @@ impl PostingListReader { } } - async fn read_positions(&self, token_id: u32) -> Result { - let positions = self.index_cache.get_or_insert_with_key(PositionKey { token_id }, || async move { + async fn read_positions( + &self, + token_id: u32, + metrics: &dyn MetricsCollector, + ) -> Result { + let result = self.index_cache.get_or_insert_with_key_hit(PositionKey { token_id }, || async move { let positions = match self.positions_layout { PositionsLayout::None => { return Err(Error::invalid_input( @@ -3853,7 +3912,12 @@ impl PostingListReader { } }; Result::Ok(Positions(positions)) - }).await?; + }).await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + let (positions, _) = result?; Ok(positions.0.clone()) } @@ -8685,12 +8749,12 @@ mod tests { // forcing the V2-only bulk metadata load. let pl_0_0 = index.partitions[0] .inverted_list - .posting_len_for_token(0) + .posting_len_for_token(0, None) .await .unwrap(); let pl_1_0 = index.partitions[1] .inverted_list - .posting_len_for_token(0) + .posting_len_for_token(0, None) .await .unwrap(); if index.partitions[0].id() == 0 { @@ -10039,7 +10103,7 @@ mod tests { assert_eq!(counter.rows_read(), 0); let (total_tokens, num_docs, dfs) = index - .bm25_stats_for_terms(&["t0".to_string()]) + .bm25_stats_for_terms(&["t0".to_string()], None) .await .unwrap(); assert_eq!(total_tokens, num_tokens as u64); @@ -10067,11 +10131,11 @@ mod tests { let (index, counter, _tmpdir) = load_counted_v2_index(100, cache.clone()).await; let terms = ["t0".to_string()]; - let first = index.bm25_stats_for_terms(&terms).await.unwrap(); + let first = index.bm25_stats_for_terms(&terms, None).await.unwrap(); assert_eq!(first, (100, 100, vec![1])); assert_eq!(counter.metadata_rows_read(), 1); - let second = index.bm25_stats_for_terms(&terms).await.unwrap(); + let second = index.bm25_stats_for_terms(&terms, None).await.unwrap(); assert_eq!(second, first); assert_eq!( counter.metadata_rows_read(), @@ -10080,6 +10144,57 @@ mod tests { ); } + /// Guards the review fix that threads `Option<&dyn MetricsCollector>` + /// through `bm25_stats_for_terms → df_for_term → posting_len_for_token → + /// posting_metadata_for_token`. Cold stats for `N` tokens on one + /// partition must record exactly `N` misses (one per `PostingMetadataKey`) + /// and zero hits; a second call with the cache warm must flip that to + /// zero misses and `N` hits. If any hop in the chain drops the collector + /// this test regresses to `0/0` on both calls. + #[tokio::test] + async fn test_bm25_stats_for_terms_records_metadata_cache_stats() { + // Keep a live `LanceCache` clone in test scope so the + // `WeakLanceCache` inside `PostingListReader` can still upgrade after + // `load_counted_v2_index` returns. Without this the weak reference + // would collapse and every `posting_metadata_for_token` call would + // silently fall through to the "cache no longer available" path, + // making both cold and warm calls report identical miss counts. + let cache = LanceCache::with_capacity(1024 * 1024); + let (index, _counter, _tmpdir) = load_counted_v2_index(100, cache.clone()).await; + assert!( + !index.partitions[0].inverted_list.is_legacy_layout(), + "this test only proves the v2 metadata boundary", + ); + + let terms = ["t0".to_string(), "t1".to_string(), "t2".to_string()]; + + let cold = LocalMetricsCollector::default(); + let cold_stats = index + .bm25_stats_for_terms(&terms, Some(&cold)) + .await + .unwrap(); + assert_eq!(cold_stats.2, vec![1, 1, 1]); + assert_eq!( + cold.index_cache_misses(), + terms.len(), + "expected one miss per (term, partition) on cold", + ); + assert_eq!(cold.index_cache_hits(), 0); + + let warm = LocalMetricsCollector::default(); + let warm_stats = index + .bm25_stats_for_terms(&terms, Some(&warm)) + .await + .unwrap(); + assert_eq!(warm_stats, cold_stats); + assert_eq!(warm.index_cache_misses(), 0); + assert_eq!( + warm.index_cache_hits(), + terms.len(), + "expected one hit per (term, partition) on warm", + ); + } + #[tokio::test] async fn test_aggregate_corpus_stats_reuses_cached_value() { let (index, _counter, _tmpdir) = load_counted_v2_index(100, LanceCache::no_cache()).await; @@ -10985,7 +11100,7 @@ mod tests { let params = Arc::new(FtsSearchParams::new().with_limit(Some(1))); let scorer = Arc::new( index - .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .bm25_base_scorer(tokens.as_ref(), params.as_ref(), None) .await .unwrap(), ); @@ -11117,7 +11232,7 @@ mod tests { assert_eq!(row_ids.len(), scores.len()); let scorer = index - .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .bm25_base_scorer(tokens.as_ref(), params.as_ref(), None) .await .unwrap(); let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); @@ -11166,7 +11281,7 @@ mod tests { assert_eq!(row_ids, vec![200]); assert_eq!(scores.len(), 1); let scorer = index - .bm25_base_scorer(tokens.as_ref(), params.as_ref()) + .bm25_base_scorer(tokens.as_ref(), params.as_ref(), None) .await .unwrap(); let expected_score = scorer.query_weight("alpha") * scorer.doc_weight(1, 1_000); diff --git a/rust/lance-index/src/scalar/ngram.rs b/rust/lance-index/src/scalar/ngram.rs index 810f100ec0c..b9373e0bf7d 100644 --- a/rust/lance-index/src/scalar/ngram.rs +++ b/rust/lance-index/src/scalar/ngram.rs @@ -249,7 +249,7 @@ impl NGramPostingListReader { row_offset: u32, metrics: &dyn MetricsCollector, ) -> Result> { - self.index_cache.get_or_insert_with_key(NGramPostingListKey { row_offset }, || async move { + let result = self.index_cache.get_or_insert_with_key_hit(NGramPostingListKey { row_offset }, || async move { metrics.record_part_load(); tracing::info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_SCALAR_PART, index_type="ngram", part_id=row_offset); let batch = self @@ -260,7 +260,12 @@ impl NGramPostingListReader { ) .await?; NGramPostingList::try_from_batch(batch, self.frag_reuse_index.clone()) - }).await + }).await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + result.map(|(v, _)| v) } } diff --git a/rust/lance-index/src/scalar/rtree.rs b/rust/lance-index/src/scalar/rtree.rs index a83f2250ac5..734bcde0ed7 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/rust/lance-index/src/scalar/rtree.rs @@ -345,15 +345,19 @@ impl RTreeIndex { while let Some(page_idx) = stack.pop() { let range = self.page_range(page_idx).await?; let is_leaf = range.start < self.metadata.num_items; - let batch = self + let result = self .index_cache - .get_or_insert_with_key(RTreeCacheKey::Page(page_idx), move || async move { + .get_or_insert_with_key_hit(RTreeCacheKey::Page(page_idx), move || async move { let batch = self.pages_reader.read_range(range, None).await?; metrics.record_part_load(); Ok(RTreeCacheValue(Arc::new(batch))) }) - .await - .map(|v| v.0.clone())?; + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + let batch = result.map(|(v, _)| v.0.clone())?; let bbox_array = extract_bounding_boxes(batch.column(0).as_ref(), batch.schema().field(0))?; @@ -381,16 +385,20 @@ impl RTreeIndex { } async fn search_null(&self, metrics: &dyn MetricsCollector) -> Result { - let batch = self + let result = self .index_cache - .get_or_insert_with_key(RTreeCacheKey::Nulls, move || async move { + .get_or_insert_with_key_hit(RTreeCacheKey::Nulls, move || async move { // Only one row let batch = self.nulls_reader.read_range(0..1, None).await?; metrics.record_part_load(); Ok(RTreeCacheValue(Arc::new(batch))) }) - .await - .map(|v| v.0.clone())?; + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + let batch = result.map(|(v, _)| v.0.clone())?; let null_map = match batch.num_rows() { 0 => RowAddrTreeMap::default(), diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index c8d8a15de55..a52ac6d9bc7 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -1199,20 +1199,27 @@ impl IVFIndex { let cache_key = IVFPartitionKey::::new(partition_id); if write_cache { - let entry = self + let result = self .index_cache - .get_or_insert_with_key(cache_key, || async { + .get_or_insert_with_key_hit(cache_key, || async { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); self.load_partition_entry(partition_id, metrics.io_stats()) .await }) - .await?; + .await; + match &result { + Ok((_, true)) => metrics.record_index_cache_hit(), + _ => metrics.record_index_cache_miss(), + } + let (entry, _) = result?; Ok(entry as Arc) } else { if let Some(part_idx) = self.index_cache.get_with_key(&cache_key).await { + metrics.record_index_cache_hit(); return Ok(part_idx); } + metrics.record_index_cache_miss(); info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); Ok(Arc::new( diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index f67ed5017c8..7d66a3633a3 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -311,6 +311,14 @@ impl MetricsCollector for FtsIndexMetrics { self.index_metrics.record_comparisons(num_comparisons); } + fn record_index_cache_hits(&self, num_hits: usize) { + self.index_metrics.record_index_cache_hits(num_hits); + } + + fn record_index_cache_misses(&self, num_misses: usize) { + self.index_metrics.record_index_cache_misses(num_misses); + } + fn record_and_candidates_seen(&self, num_candidates: usize) { self.and_candidates_seen.add(num_candidates); } @@ -668,9 +676,14 @@ impl ExecutionPlan for MatchQueryExec { None => { let scorer_start = std::time::Instant::now(); let scorer = Arc::new( - build_global_bm25_scorer(&indices, &tokens, ¶ms) - .boxed() - .await?, + build_global_bm25_scorer( + &indices, + &tokens, + ¶ms, + Some(metrics.as_ref()), + ) + .boxed() + .await?, ); metrics.record_scorer_build(scorer_start.elapsed()); scorer @@ -1220,6 +1233,7 @@ impl ExecutionPlan for FlatMatchQueryExec { &indices, &query_tokens, &FtsSearchParams::new(), + Some(metrics.as_ref()), ) .boxed() .await?; @@ -1570,9 +1584,14 @@ impl ExecutionPlan for PhraseQueryExec { None => { let scorer_start = std::time::Instant::now(); let scorer = Arc::new( - build_global_bm25_scorer(&indices, &tokens, ¶ms) - .boxed() - .await?, + build_global_bm25_scorer( + &indices, + &tokens, + ¶ms, + Some(metrics.as_ref()), + ) + .boxed() + .await?, ); metrics.record_scorer_build(scorer_start.elapsed()); scorer @@ -3081,7 +3100,7 @@ mod tests { let mut tokenizer = indices[0].tokenizer(); let tokens = collect_query_tokens(&query.terms, &mut tokenizer); let global_scorer = Arc::new( - build_global_bm25_scorer(&indices, &tokens, &search_params) + build_global_bm25_scorer(&indices, &tokens, &search_params, None) .await .unwrap(), ); diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index 1af0bc3f4ef..f4e4b99512f 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -2,8 +2,9 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use lance_datafusion::utils::{ - BYTES_READ_METRIC, ExecutionPlanMetricsSetExt, INDEX_COMPARISONS_METRIC, INDICES_LOADED_METRIC, - IOPS_METRIC, PARTS_LOADED_METRIC, REQUESTS_METRIC, + BYTES_READ_METRIC, ExecutionPlanMetricsSetExt, INDEX_CACHE_HITS_METRIC, + INDEX_CACHE_MISSES_METRIC, INDEX_COMPARISONS_METRIC, INDICES_LOADED_METRIC, IOPS_METRIC, + PARTS_LOADED_METRIC, REQUESTS_METRIC, }; use lance_index::metrics::MetricsCollector; use lance_io::scheduler::{IoStats, ScanScheduler, ScanStats}; @@ -517,6 +518,8 @@ pub struct IndexMetrics { indices_loaded: Count, parts_loaded: Count, index_comparisons: Count, + index_cache_hits: Count, + index_cache_misses: Count, /// Per-query sink that accumulates exact index-file I/O as partitions are /// loaded from storage. Shared by all clones of this `IndexMetrics`, so /// concurrent partition loads all funnel into the same counters. Published @@ -531,6 +534,8 @@ impl IndexMetrics { indices_loaded: metrics.new_count(INDICES_LOADED_METRIC, partition), parts_loaded: metrics.new_count(PARTS_LOADED_METRIC, partition), index_comparisons: metrics.new_count(INDEX_COMPARISONS_METRIC, partition), + index_cache_hits: metrics.new_count(INDEX_CACHE_HITS_METRIC, partition), + index_cache_misses: metrics.new_count(INDEX_CACHE_MISSES_METRIC, partition), io_stats: IoStats::new(), io_metrics: IoMetrics::new(metrics, partition), } @@ -555,6 +560,12 @@ impl MetricsCollector for IndexMetrics { fn record_comparisons(&self, num_comparisons: usize) { self.index_comparisons.add(num_comparisons); } + fn record_index_cache_hits(&self, num_hits: usize) { + self.index_cache_hits.add(num_hits); + } + fn record_index_cache_misses(&self, num_misses: usize) { + self.index_cache_misses.add(num_misses); + } fn io_stats(&self) -> Option { Some(self.io_stats.clone()) }