From 7d3c62369ad074a4cd480b02e75089f885356c96 Mon Sep 17 00:00:00 2001 From: Walter Woodall Date: Mon, 27 Jul 2026 13:58:51 -0700 Subject: [PATCH 1/2] add: Serialize to ProbeStats --- src/vector/backend.rs | 75 ++++++++++++++++++++++++++++++++++++++--- src/vector/ivf/graph.rs | 4 +-- src/vector/ivf/index.rs | 2 +- 3 files changed, 74 insertions(+), 7 deletions(-) diff --git a/src/vector/backend.rs b/src/vector/backend.rs index bd27ac191c..572d7d8e67 100644 --- a/src/vector/backend.rs +++ b/src/vector/backend.rs @@ -141,7 +141,7 @@ impl VectorBackend { } /// How the probe loop stopped. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, serde::Serialize)] pub enum ProbeTermination { /// The filter-effective probe budget reached `max_probe_count` — the /// probe ceiling. @@ -158,7 +158,7 @@ pub enum ProbeTermination { /// touched. Returned by [`VectorBackend::top_n`] alongside the hits. /// The flat/exact path fills only `exact_rows_read`; every other field /// is IVF-probe-only. -#[derive(Debug, Default)] +#[derive(Debug, Default, serde::Serialize)] pub struct ProbeStats { /// Clusters visited by the probe loop, in probe order. A cluster /// appears here once we've passed the stop-condition gate for it, @@ -625,8 +625,9 @@ mod tests { use crate::schema::{IndexRecordOption, Schema, Term, STORED, STRING}; use crate::vector::tests::{exhaustive_params, TestVectorIndex}; use crate::vector::{ - IvfCentroids, IvfClusterer, IvfMatrix, IvfMergeSettings, IvfVectors, VectorClusterStats, - VectorDType, VectorInfo, VectorOptions, VectorStorageFormat, + IvfCentroids, IvfClusterer, IvfMatrix, IvfMergeSettings, IvfVectors, + NeighborhoodGraphSearchMetrics, SearchTerminationReason, VectorClusterStats, VectorDType, + VectorInfo, VectorOptions, VectorStorageFormat, }; use crate::{Index, IndexWriter, TantivyDocument}; @@ -2341,6 +2342,72 @@ mod tests { Ok(()) } + /// `ProbeStats` (and nested routing / optional graph metrics) round-trip + /// through `serde_json` with the field names callers rely on. + #[test] + fn probe_stats_serializes_to_json() { + let stats = ProbeStats { + probed_clusters: vec![2, 5], + candidates_scored: 10, + vectors_visited: 20, + pruned_filter: 4, + pruned_dead: 3, + pruned_seen: 3, + postings_row: 1, + postings_skipped: 1, + exact_rows_read: 0, + routing: IvfSearchMetrics { + visited_count: 7, + graph: Some(NeighborhoodGraphSearchMetrics { + visited_count: 7, + expanded_count: 4, + edges_scanned: 12, + evictions: 1, + result_count: 3, + termination_reason: SearchTerminationReason::SearchConverged, + }), + }, + min_candidates: 5, + termination: ProbeTermination::Gate, + }; + + let value = serde_json::to_value(&stats).expect("ProbeStats should serialize to JSON"); + assert_eq!( + value, + serde_json::json!({ + "probed_clusters": [2, 5], + "candidates_scored": 10, + "vectors_visited": 20, + "pruned_filter": 4, + "pruned_dead": 3, + "pruned_seen": 3, + "postings_row": 1, + "postings_skipped": 1, + "exact_rows_read": 0, + "routing": { + "visited_count": 7, + "graph": { + "visited_count": 7, + "expanded_count": 4, + "edges_scanned": 12, + "evictions": 1, + "result_count": 3, + "termination_reason": "SearchConverged" + } + }, + "min_candidates": 5, + "termination": "Gate" + }) + ); + + // Exact routing leaves `graph` unset — still must serialize as null. + let mut exact_routing = stats; + exact_routing.routing.graph = None; + let exact_value = + serde_json::to_value(&exact_routing).expect("ProbeStats should serialize to JSON"); + assert_eq!(exact_value["routing"]["graph"], serde_json::Value::Null); + } + // ============================================================ // Filter-aware posting fetches. // diff --git a/src/vector/ivf/graph.rs b/src/vector/ivf/graph.rs index cf98fbfca9..3100eecc0e 100644 --- a/src/vector/ivf/graph.rs +++ b/src/vector/ivf/graph.rs @@ -378,7 +378,7 @@ impl EdgeListMut<'_> { } /// Why a [`RelativeNeighborhoodGraph::search`] stopped expanding. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)] pub enum SearchTerminationReason { /// The best unexpanded candidate could not beat the worst kept result of /// a full beam: the search converged. @@ -390,7 +390,7 @@ pub enum SearchTerminationReason { /// Per-query cost and convergence counters returned by /// [`RelativeNeighborhoodGraph::search`]. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, serde::Serialize)] pub struct NeighborhoodGraphSearchMetrics { /// Nodes visited — and therefore scored — by the query; the search's /// navigation cost. diff --git a/src/vector/ivf/index.rs b/src/vector/ivf/index.rs index 15525ec381..2942f7e04b 100644 --- a/src/vector/ivf/index.rs +++ b/src/vector/ivf/index.rs @@ -331,7 +331,7 @@ impl Iterator for ClusterRanking<'_> { /// [`ClusterRanking::metrics`] snapshot): how many centroids were scored to /// pick the probe order, and — when routing went through the centroid RNG — /// the beam search's full [`NeighborhoodGraphSearchMetrics`]. -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Copy, Debug, Default, serde::Serialize)] pub struct IvfSearchMetrics { /// Centroids scored to route the query (the navigation cost): /// `num_centroids` on the exact path, the beam-visited count when routed From 4421729757ac8634dc50127e4eb7daad385510f8 Mon Sep 17 00:00:00 2001 From: Walter Woodall Date: Mon, 27 Jul 2026 15:58:09 -0700 Subject: [PATCH 2/2] fix: remove probed_clusters from ProbeStats The probed_clusters attributed resulted in unusable JSON blobs since it logged the id of each probed cluster. This was useful for one off debugging but not for production level observability / metrics. --- src/vector/backend.rs | 172 +++++++++++----------------------------- src/vector/collector.rs | 4 +- 2 files changed, 47 insertions(+), 129 deletions(-) diff --git a/src/vector/backend.rs b/src/vector/backend.rs index 572d7d8e67..6f95a41b7a 100644 --- a/src/vector/backend.rs +++ b/src/vector/backend.rs @@ -153,17 +153,12 @@ pub enum ProbeTermination { Exhausted, } -/// Per-segment probe-loop instrumentation: which clusters were probed -/// (in probe order) and a prune breakdown of every doc the inner loop -/// touched. Returned by [`VectorBackend::top_n`] alongside the hits. -/// The flat/exact path fills only `exact_rows_read`; every other field -/// is IVF-probe-only. +/// Per-segment probe-loop instrumentation: a prune breakdown of every +/// doc the inner loop touched, plus posting-fetch counters. Returned by +/// [`VectorBackend::top_n`] alongside the hits. The flat/exact path fills +/// only `exact_rows_read`; every other field is IVF-probe-only. #[derive(Debug, Default, serde::Serialize)] pub struct ProbeStats { - /// Clusters visited by the probe loop, in probe order. A cluster - /// appears here once we've passed the stop-condition gate for it, - /// regardless of whether its doc-ids slice ends up empty. - pub probed_clusters: Vec, /// Docs that passed filter + alive + seen and were scored against the /// query. This stays the "scored" bucket and equals the final survivor /// `candidates`, so starvation is just `candidates_scored < min_candidates`. @@ -185,7 +180,7 @@ pub struct ProbeStats { /// `filter → alive → seen` pre-pass left zero survivors (fully /// filtered / dead / already-seen, or the cluster is empty). The two /// `postings_*` counters partition the probed clusters: - /// `postings_row + postings_skipped == probed_clusters.len()`. + /// [`clusters_probed`](Self::clusters_probed) `== postings_row + postings_skipped`. pub postings_skipped: usize, /// Flat/exact-path stride-sized row reads — one per survivor scored. /// Filled only by the exact (non-IVF) path. @@ -202,6 +197,15 @@ pub struct ProbeStats { pub termination: ProbeTermination, } +impl ProbeStats { + /// Clusters the probe loop visited — each either fetched survivors + /// (`postings_row`) or skipped (`postings_skipped`). + #[inline] + pub fn clusters_probed(&self) -> usize { + self.postings_row + self.postings_skipped + } +} + /// Floor a probed cluster charges the ceiling even when the filter skips /// all its rows (the gate pre-pass still scans them). A cluster bills /// `SKIPPED_CLUSTER_COST + (1 - SKIPPED_CLUSTER_COST) * pass_fraction`: @@ -375,10 +379,6 @@ impl VectorBackend { } let cluster = cluster as usize; - // Record the probe before doing any work, so even an empty - // cluster counts as "probed". - stats.probed_clusters.push(cluster); - let rows = index.cluster_range(cluster); let num_rows = rows.len(); @@ -1855,15 +1855,15 @@ mod tests { assert!(hits.is_empty()); // Short-circuit fires before the probe loop, so no clusters // visited and no candidates scored. - assert!(stats.probed_clusters.is_empty()); + assert_eq!(stats.clusters_probed(), 0); assert_eq!(stats.candidates_scored, 0); Ok(()) } - /// Smoke for the instrumented seam: probed_clusters is non-empty, - /// every entry is < num_centroids, and candidates_scored ≤ total - /// docs in the inspected segment. Exhaustive params on a 9-centroid - /// segment visit all 9. + /// Smoke for the instrumented seam: every centroid is probed under + /// exhaustive params, and candidates_scored ≤ total docs in the + /// inspected segment. Exhaustive params on a 9-centroid segment + /// visit all 9. #[test] fn ivf_top_n_collects_probe_stats() -> crate::Result<()> { let index = TestVectorIndex::builder(VectorDType::F32) @@ -1877,10 +1877,7 @@ mod tests { 4, exhaustive_params(DEFAULT_NUM_CENTROIDS), )?; - assert_eq!(stats.probed_clusters.len(), DEFAULT_NUM_CENTROIDS); - for &c in &stats.probed_clusters { - assert!(c < DEFAULT_NUM_CENTROIDS, "probed cluster {c} out of range"); - } + assert_eq!(stats.clusters_probed(), DEFAULT_NUM_CENTROIDS); // The first segment has docs distributed across all 9 clusters; // candidates_scored equals the segment's doc count under // exhaustive probe + AllQuery. @@ -1942,7 +1939,7 @@ mod tests { let (_, stats) = run_top_n(&index, embed_field, vec![10.0, 10.0], 3, params)?; assert_eq!(stats.termination, ProbeTermination::Ceiling); // Stopped at exactly the cap, short of the ranked list. - assert_eq!(stats.probed_clusters.len(), 1); + assert_eq!(stats.clusters_probed(), 1); assert_eq!(stats.routing.visited_count, centroids.len()); assert_eq!( stats.vectors_visited, @@ -1983,7 +1980,7 @@ mod tests { AdaptiveProbeParams::default(), )?; assert_eq!(hits, expected, "linear fallback must match the oracle"); - assert_eq!(stats.probed_clusters, vec![0], "one cluster, one probe"); + assert_eq!(stats.clusters_probed(), 1, "one cluster, one probe"); assert_eq!(stats.routing.visited_count, 1); Ok(()) } @@ -2038,9 +2035,9 @@ mod tests { let (hits, stats) = run_top_n(&index, embed_field, query.to_vec(), k, params.clone())?; assert_eq!(hits, expected, "routed top-{k} near centroid {ord}"); assert!( - stats.probed_clusters.len() <= 2, - "cap 2 must bound the probes, got {:?}", - stats.probed_clusters + stats.clusters_probed() <= 2, + "cap 2 must bound the probes, got {}", + stats.clusters_probed() ); assert!( stats.routing.visited_count <= centroids.len(), @@ -2164,20 +2161,12 @@ mod tests { )?; assert_eq!(stats.termination, ProbeTermination::Ceiling); - let searcher = index.index.reader()?.searcher(); - let sizes = searcher.segment_readers()[0] - .vector_index(index.embedding_field())? - .cluster_sizes() - .expect("ivf segment exposes cluster sizes"); - let non_empty_probed = stats - .probed_clusters - .iter() - .filter(|&&c| sizes[c] > 0) - .count(); + // With AllQuery every non-empty probed cluster fetches survivors + // (`postings_row`); empty ones skip. The filter-effective ceiling + // of 2 therefore binds at exactly 2 fetches. assert_eq!( - non_empty_probed, 2, - "cap ⇒ exactly 2 non-empty (filter-effective) clusters probed, got {:?}", - stats.probed_clusters, + stats.postings_row, 2, + "cap ⇒ exactly 2 non-empty (filter-effective) clusters probed, got {stats:?}" ); Ok(()) } @@ -2253,10 +2242,9 @@ mod tests { let (_, stats) = run_top_n(&index, embed_field, vec![1.0, 0.3], 1, params)?; assert_eq!(stats.termination, ProbeTermination::Gate); assert_eq!( - stats.probed_clusters.len(), + stats.clusters_probed(), 1, - "gate must stop before the far angular cluster ({:?})", - stats.probed_clusters, + "gate must stop before the far angular cluster ({stats:?})", ); Ok(()) } @@ -2288,56 +2276,10 @@ mod tests { params, )?; assert!( - stats.probed_clusters.len() < DEFAULT_NUM_CENTROIDS, + stats.clusters_probed() < DEFAULT_NUM_CENTROIDS, "default-params pruning should visit strictly fewer than {DEFAULT_NUM_CENTROIDS} \ - clusters; got {} ({:?})", - stats.probed_clusters.len(), - stats.probed_clusters, - ); - Ok(()) - } - - /// Structural invariants on the probe stats themselves — - /// independent of any specific stop-condition behavior. - /// - all probed indices live in [0, num_centroids) - /// - no duplicates (a cluster is probed at most once) - /// - the first probed cluster is the centroid nearest the query - #[test] - fn probe_stats_probed_clusters_validity() -> crate::Result<()> { - let index = TestVectorIndex::builder(VectorDType::F32) - .vector_storage_format(VectorStorageFormat::Ivf) - .build()?; - let query = [9.0_f32, 0.5]; - let (_, stats) = run_top_n( - &index.index, - index.embedding_field(), - query.to_vec(), - 2, - exhaustive_params(DEFAULT_NUM_CENTROIDS), - )?; - - for &c in &stats.probed_clusters { - assert!( - c < DEFAULT_NUM_CENTROIDS, - "probed cluster {c} out of range (num_centroids={DEFAULT_NUM_CENTROIDS})", - ); - } - let unique: std::collections::HashSet = - stats.probed_clusters.iter().copied().collect(); - assert_eq!( - unique.len(), - stats.probed_clusters.len(), - "duplicate probed cluster: {:?}", - stats.probed_clusters, - ); - - let nearest = nearest_centroid_to(&query); - assert_eq!( - stats.probed_clusters.first().copied(), - Some(nearest), - "first probed should be the centroid nearest the query; nearest = {nearest}, \ - probed_clusters = {:?}", - stats.probed_clusters, + clusters; got {} ({stats:?})", + stats.clusters_probed(), ); Ok(()) } @@ -2347,7 +2289,6 @@ mod tests { #[test] fn probe_stats_serializes_to_json() { let stats = ProbeStats { - probed_clusters: vec![2, 5], candidates_scored: 10, vectors_visited: 20, pruned_filter: 4, @@ -2375,7 +2316,6 @@ mod tests { assert_eq!( value, serde_json::json!({ - "probed_clusters": [2, 5], "candidates_scored": 10, "vectors_visited": 20, "pruned_filter": 4, @@ -2399,6 +2339,7 @@ mod tests { "termination": "Gate" }) ); + assert_eq!(stats.clusters_probed(), 2); // Exact routing leaves `graph` unset — still must serialize as null. let mut exact_routing = stats; @@ -2466,20 +2407,13 @@ mod tests { backend.top_n(weight, segment_reader, k) } - /// The two partition identities every scan must uphold: each touched - /// row lands in exactly one prune bucket, and each probed cluster - /// either fetches its survivors or skips. + /// Every touched row lands in exactly one prune bucket. fn assert_stats_identities(stats: &ProbeStats) { assert_eq!( stats.vectors_visited, stats.pruned_filter + stats.pruned_dead + stats.pruned_seen + stats.candidates_scored, "visited must equal filter+dead+seen+scored ({stats:?})" ); - assert_eq!( - stats.probed_clusters.len(), - stats.postings_row + stats.postings_skipped, - "probed clusters must partition into fetched+skipped ({stats:?})" - ); } /// Build a single-segment FLAT index (one commit, never merged past @@ -2579,7 +2513,7 @@ mod tests { assert_eq!(stats.postings_row, 0, "0%: no fetches"); assert_eq!( stats.postings_skipped, - stats.probed_clusters.len(), + stats.clusters_probed(), "0%: every probed cluster skips its fetch" ); } @@ -2590,7 +2524,7 @@ mod tests { assert_eq!(stats.postings_row, 1, "{stats:?}"); assert_eq!( stats.postings_skipped, - stats.probed_clusters.len() - 1, + stats.clusters_probed() - 1, "{stats:?}" ); } @@ -2635,7 +2569,7 @@ mod tests { ); // Only the first-probed cell fetches anything. assert_eq!(stats.postings_row, 1, "{stats:?}"); - assert_eq!(stats.postings_skipped, stats.probed_clusters.len() - 1); + assert_eq!(stats.postings_skipped, stats.clusters_probed() - 1); assert_stats_identities(&stats); Ok(()) } @@ -2697,10 +2631,9 @@ mod tests { Ok(()) } - /// An empty cluster still counts as probed — the probe is recorded - /// before any work — and takes the skip path: `postings_skipped` - /// increments, nothing is fetched, and the visited/prune counters - /// don't move. + /// An empty cluster still counts as probed — the loop visits it and + /// takes the skip path: `postings_skipped` increments, nothing is + /// fetched, and the visited/prune counters don't move. #[test] fn empty_cluster_probed_but_fetch_skipped() -> crate::Result<()> { // No doc is nearest to the third centroid, so its cluster is @@ -2739,14 +2672,10 @@ mod tests { )?; assert_eq!(hits.len(), 8); assert_eq!( - stats.probed_clusters.len(), + stats.clusters_probed(), centroids.len(), "the empty cluster still counts as probed: {stats:?}" ); - assert!( - stats.probed_clusters.contains(&2), - "the empty cluster is in the probe list: {stats:?}" - ); assert_eq!( stats.postings_skipped, 1, "the empty cluster fetched nothing: {stats:?}" @@ -2794,7 +2723,7 @@ mod tests { // fields must stay zeroed. assert_eq!(stats.vectors_visited, 0, "{pct}%: {stats:?}"); assert_eq!(stats.candidates_scored, 0, "{pct}%: {stats:?}"); - assert!(stats.probed_clusters.is_empty(), "{pct}%: {stats:?}"); + assert_eq!(stats.clusters_probed(), 0, "{pct}%: {stats:?}"); assert_eq!( stats.exact_rows_read, admitted, "{pct}%: one row read per survivor" @@ -2945,15 +2874,4 @@ mod tests { fn grid2d_first_centroid() -> [f32; 2] { [0.0, 0.0] } - - /// L2-nearest centroid index for a query against the shared - /// fixture's default 3×3 grid centroids. - fn nearest_centroid_to(query: &[f32; 2]) -> usize { - // Match the grid in `crate::vector::tests::grid2d::centroids()`: - // origin=(0,0), 3×3, gap=3.0, row-major. - let centroids: Vec<[f32; 2]> = (0..3) - .flat_map(|row| (0..3).map(move |col| [col as f32 * 3.0, row as f32 * 3.0])) - .collect(); - nearest_centroid(*query, ¢roids) - } } diff --git a/src/vector/collector.rs b/src/vector/collector.rs index 160cd37008..167c0a69f1 100644 --- a/src/vector/collector.rs +++ b/src/vector/collector.rs @@ -78,8 +78,8 @@ pub struct VectorSimilarityFruit { pub results: Vec<(Score, DocAddress)>, /// One [`ProbeStats`] per collected segment, in segment-ordinal order /// after [`Collector::merge_fruits`]. The counter fields are summable - /// across segments; `probed_clusters`, `min_candidates`, and - /// `termination` only carry per-segment meaning. + /// across segments; `min_candidates` and `termination` only carry + /// per-segment meaning. pub stats: Vec, }