seaweed-vector: probe policy for indexed queries (ANN phase 1) (#12412)

* seaweed-vector: build and extend the vector index weed_lance_worker::vector_index builds an IVF_FLAT index in the metric the table was created with, sizes partitions from the row count, and extends an index with a new segment or a retrain; the core crate re-exports it and adds the probes a query spends. Tests cover coverage after build, append, retrain and compaction. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector-cli: index build and index optimize Logs go to stderr so stdout stays the JSON a caller parses. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: PERF_INDEX builds the index before the timed queries The cluster fixture hands out its catalog URL, the Makefile builds weed-vector-cli, and the perf suite builds the index after ingest and reports the build time and segment count. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: probe a tenth of the index's partitions, more under a filter A query that names no probe count gets one from the index's partition count. Lance searches at least that many partitions, and under a prefilter keeps going until it has k rows; without one it stops there. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: --nprobes overrides the probe policy Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: test that a build takes the table's metric Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: make clean removes weed-vector-cli Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: measure recall with vectors the table does not hold Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: keep the read cache's doc comment on with_read_cache Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: optimize only the indices on the vector column Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: resolve the query metric with index::distance_metric Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * seaweed-vector: test the L2 metric under the name CreateIndex stores Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * weed-worker: open a table inside a caller's Lance session Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: keep opened indices between queries in one session Every open used a fresh Lance session, so the index and metadata caches died with the request and each query loaded the IVF index again. --index-cache-mb sizes the shared cache, 512 MiB by default. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: PERF_CLUSTERS draws the perf suite's vectors around shared centres Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: test that compaction drops the deleted key from tied hits Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * seaweed-vector: test that an append-only optimize serves the new row Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: record PERF_NOISE in the perf report config Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * seaweed-vector: a filtered query widens to four times the probes, not without bound With no upper bound Lance prepares every partition for the late search even when the minimum probes already hold k rows, and a filtered query through the index cost as much as the exact scan. A query that names its own probe count gets exactly that count. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: report recall through an index without asserting it Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: count a hit tying the k-th best score toward recall Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: allow float32 rounding at the recall cutoff Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: count a repeated hit key once toward recall Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: credit cutoff ties only up to the slots the top-k leaves Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test/s3vectors: assert every filtered hit passes the filter, indexed or not Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * weed-vector: forget the shared session when a table was dropped and created again The session keys index lists and transactions by dataset version, and a table created again at the same location restarts its versions. An open that finds an older version than the newest seen at that URI, or the same version committed later, replaces the session and opens again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * weed-vector: test an open after another engine created the table again Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * weed-vector: replace the session in the same lock that records a recreated table A second open of the recreated table could see the new version recorded and the old session still shared, and keep the dropped table's index list. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * seaweed-vector: size default probes from the largest index on the column A sharded build gives the column one index name per shard. The first name alone could be a small shard, and its probe count would then be spent on a larger one. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: skip the probe lookup when a query is given no segments Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: open every table in its own Lance session again One session shared by every open keeps a table's index loaded between queries, but Lance keys what it caches by location and version, so a table dropped and created again at the same location could be served the old table's index list. Telling the two apart from version and commit time needs a map that only grows, a retry loop and clocks that agree across engines. A private session per open, as before, has no such state. Measured at 50k x 384 through an index, eight clients: 164 queries/s with the shared session, 137 without; one client is unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: reject --nprobes 0 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: read the partition count from the opened index Lance's index statistics took 3.6 ms a query on a 200k-row table with 447 partitions. Opening the index and asking it takes 0.4 ms, and the scan that follows finds the index already open. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-11-05 11:37:44 +01:00 · 2026-11-04 21:07:13 -08:00
1 parent f0f0ba85ae
commit b5951c5c7c
@@ -1,18 +1,60 @@
 //! Sizing the vector index an S3 Vectors table searches through, and the
 //! probes a query spends on it. Building lives in the worker crate.

+use anyhow::Result;
+use lance::Dataset;
+use lance::index::{DatasetIndexExt, DatasetIndexInternalExt};
+
 pub use weed_lance_worker::vector_index::{
    INDEX_NAME, MAX_PARTITIONS, build, distance_metric, index_names, optimize, partitions_for,
 };

 pub const MIN_PROBES: usize = 8;

-/// Partitions one query probes: a tenth of them, at least [`MIN_PROBES`], never
-/// more than exist.
-pub fn probes_for(partitions: usize) -> usize {
-    MIN_PROBES
-        .max(partitions.div_ceil(10))
-        .min(partitions.max(1))
+/// How many partitions a query searches: `min` always, and under a prefilter
+/// up to `max` until it has k rows.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct Probes {
+    pub min: usize,
+    pub max: usize,
+}
+
+/// Probes for an index of `partitions`: a tenth of them, at least
+/// [`MIN_PROBES`], never more than exist; a filtered query may widen to four
+/// times that.
+pub fn probes_for(partitions: usize) -> Probes {
+    let partitions = partitions.max(1);
+    let min = MIN_PROBES.max(partitions.div_ceil(10)).min(partitions);
+    Probes {
+        min,
+        max: (min * 4).min(partitions),
+    }
+}
+
+/// The probes a query of `vector_column` uses when it names none, or `None`
+/// without an index. The count comes from the most partitions any index on
+/// the column has, whichever of them the scan ends up searching.
+/// It is read from the opened index, which the scan then finds loaded;
+/// Lance's index statistics cost several milliseconds a query.
+pub async fn default_probes(dataset: &Dataset, vector_column: &str) -> Result<Option<Probes>> {
+    let Some(field) = dataset.schema().field(vector_column) else {
+        return Ok(None);
+    };
+    let mut partitions = None;
+    for index in dataset.load_indices().await?.iter() {
+        if index.fields != [field.id] {
+            continue;
+        }
+        let opened = dataset
+            .open_vector_index(
+                vector_column,
+                &index.uuid,
+                &lance_index::metrics::NoOpMetricsCollector,
+            )
+            .await?;
+        partitions = Some(opened.total_partitions().max(partitions.unwrap_or(0)));
+    }
+    Ok(partitions.map(probes_for))
 }

 #[cfg(test)]
@@ -29,11 +71,17 @@ mod tests {

    #[test]
    fn probes_are_a_tenth_with_a_floor() {
-        assert_eq!(probes_for(1), 1);
-        assert_eq!(probes_for(8), 8);
-        assert_eq!(probes_for(64), 8);
-        assert_eq!(probes_for(224), 23);
-        assert_eq!(probes_for(4096), 410);
+        assert_eq!(probes_for(1), Probes { min: 1, max: 1 });
+        assert_eq!(probes_for(8), Probes { min: 8, max: 8 });
+        assert_eq!(probes_for(64), Probes { min: 8, max: 32 });
+        assert_eq!(probes_for(224), Probes { min: 23, max: 92 });
+        assert_eq!(
+            probes_for(4096),
+            Probes {
+                min: 410,
+                max: 1640
+            }
+        );
    }

    fn exact(vector: &[f32], k: usize) -> Query {
@@ -73,6 +121,107 @@ mod tests {
        }
    }

+    /// Pseudo-random rows, so partitions are not degenerate; a different seed
+    /// gives vectors the table does not hold.
+    fn scattered(rows: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
+        let mut state: u64 = seed;
+        (0..rows)
+            .map(|_| {
+                (0..dim)
+                    .map(|_| {
+                        state = state
+                            .wrapping_mul(6364136223846793005)
+                            .wrapping_add(1442695040888963407);
+                        ((state >> 40) as f32) / (1u64 << 24) as f32
+                    })
+                    .collect()
+            })
+            .collect()
+    }
+
+    #[tokio::test]
+    async fn default_probes_follow_the_index_and_find_the_neighbours() {
+        let dir = TempDir::new().unwrap();
+        let rows = scattered(4000, 8, 0x9E3779B97F4A7C15);
+        let tags: Vec<&str> = (0..4000)
+            .map(|i| if i % 100 == 0 { "rare" } else { "common" })
+            .collect();
+        let mut dataset = testing::write_dataset(&dir, &rows, &tags).await;
+        let table = VectorTable::resolve(&dataset, None, None).unwrap();
+        assert_eq!(default_probes(&dataset, "vector").await.unwrap(), None);
+        build(&mut dataset, &table.vector_column, INDEX_NAME, 64, false)
+            .await
+            .unwrap();
+        assert_eq!(
+            default_probes(&dataset, "vector").await.unwrap(),
+            Some(probes_for(64))
+        );
+
+        let mut found = 0;
+        let total = 20 * 10;
+        for vector in &scattered(20, 8, 0xD1B54A32D192ED03) {
+            let q = exact(vector, 10);
+            let flat: std::collections::HashSet<u64> = query::flat(&dataset, &table, &q)
+                .await
+                .unwrap()
+                .iter()
+                .map(|h| h.row_id)
+                .collect();
+            let ann = query::ann(&dataset, &table, &q).await.unwrap();
+            assert_eq!(ann.len(), 10);
+            found += ann.iter().filter(|h| flat.contains(&h.row_id)).count();
+        }
+        assert!(found * 10 >= total * 9, "recall {found}/{total}");
+
+        // Eight probes do not hold ten of the 1% rare rows; a filtered query
+        // widens past the minimum until it has them.
+        let filtered = query::ann(
+            &dataset,
+            &table,
+            &Query {
+                filter: Some("tag = 'rare'".to_string()),
+                ..exact(&rows[0], 10)
+            },
+        )
+        .await
+        .unwrap();
+        assert_eq!(filtered.len(), 10);
+    }
+
+    #[tokio::test]
+    async fn default_probes_take_the_largest_index() {
+        use lance::index::DatasetIndexExt;
+        use lance::index::vector::VectorIndexParams;
+        use lance_index::IndexType;
+        let dir = TempDir::new().unwrap();
+        let mut dataset =
+            testing::write_dataset(&dir, &scattered(1000, 8, 0x9E3779B97F4A7C15), &["a"; 1000])
+                .await;
+        testing::append_dataset(
+            &mut dataset,
+            &scattered(1000, 8, 0xD1B54A32D192ED03),
+            &["b"; 1000],
+        )
+        .await;
+        for (name, fragment, partitions) in [("a-small", 0u32, 2), ("z-large", 1u32, 100)] {
+            dataset
+                .create_index_builder(
+                    &["vector"],
+                    IndexType::Vector,
+                    &VectorIndexParams::ivf_flat(partitions, DistanceType::L2),
+                )
+                .name(name.to_string())
+                .fragments(vec![fragment])
+                .await
+                .unwrap();
+            dataset.checkout_latest().await.unwrap();
+        }
+        assert_eq!(
+            default_probes(&dataset, "vector").await.unwrap(),
+            Some(probes_for(100))
+        );
+    }
+
    #[tokio::test]
    async fn build_covers_every_fragment() {
        let dir = TempDir::new().unwrap();
@@ -120,8 +120,28 @@ async fn scanner(
        crate::filter::apply(&mut scanner, filter)?;
        scanner.prefilter(!query.postfilter);
    }
-    if let Some(nprobes) = query.nprobes {
-        scanner.nprobes(nprobes);
+    // An explicitly empty segment list searches no index, so there is
+    // nothing to size probes for.
+    let no_segments = query.index_segments.as_ref().is_some_and(|s| s.is_empty());
+    if use_index && !no_segments {
+        let probes = match query.nprobes {
+            Some(n) => Some(crate::index::Probes { min: n, max: n }),
+            None => crate::index::default_probes(dataset, &table.vector_column).await?,
+        };
+        if let Some(probes) = probes {
+            // Lance searches `min` partitions, and under a prefilter keeps
+            // going up to `max` until it has k rows. An open bound would
+            // make it prepare every partition for that, which costs as much
+            // as the exact scan it is meant to avoid.
+            scanner.minimum_nprobes(probes.min);
+            scanner.maximum_nprobes(probes.max);
+            tracing::debug!(
+                min = probes.min,
+                max = probes.max,
+                filtered = query.filter.is_some(),
+                "ann scan"
+            );
+        }
    }
    let segments = query.index_segments.as_deref().unwrap_or(&[]);
    if use_index && !segments.is_empty() {
@@ -36,6 +36,10 @@ struct Args {
    /// Memory for caching reads of write-once table files, in MiB; 0 disables.
    #[arg(long, env = "WEED_VECTOR_READ_CACHE_MB", default_value_t = 512)]
    read_cache_mb: u64,
+    /// Partitions each query probes in a table's vector index. Default: a
+    /// tenth of the index's partitions, at least 8; more under a filter.
+    #[arg(long, env = "WEED_VECTOR_NPROBES")]
+    nprobes: Option<std::num::NonZeroUsize>,
    /// The `jwt.filer_signing.key` of the cluster. The Lance namespace vends
    /// credentials that write index data only to a caller holding it.
    #[arg(long, env = "WEED_FILER_SIGNING_KEY")]
@@ -154,7 +158,8 @@ async fn main() -> Result<()> {
    let client = NamespaceClient::new(args.namespace.clone())
        .with_fallback(&fallback)
        .with_signing_key(signing_key.as_ref());
-    let mut service = VectorServer::new(client, fallback);
+    let mut service = VectorServer::new(client, fallback)
+        .with_nprobes(args.nprobes.map(std::num::NonZeroUsize::get));
    if args.read_cache_mb > 0 {
        service = service.with_read_cache(ReadCache::new(args.read_cache_mb << 20));
    }
@@ -192,3 +197,15 @@ async fn main() -> Result<()> {
    }
    Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn nprobes_must_be_positive() {
+        assert!(Args::try_parse_from(["weed-vector", "--nprobes", "0"]).is_err());
+        let args = Args::try_parse_from(["weed-vector", "--nprobes", "12"]).unwrap();
+        assert_eq!(args.nprobes.map(std::num::NonZeroUsize::get), Some(12));
+    }
+}
@@ -374,6 +374,7 @@ pub struct VectorServer {
    client: NamespaceClient,
    fallback: FallbackOptions,
    read_cache: Option<Arc<dyn WrappingObjectStore>>,
+    nprobes: Option<usize>,
 }

 impl VectorServer {
@@ -382,9 +383,16 @@ impl VectorServer {
            client,
            fallback,
            read_cache: None,
+            nprobes: None,
        }
    }

+    /// Partitions every query probes, instead of the index's own default.
+    pub fn with_nprobes(mut self, nprobes: Option<usize>) -> Self {
+        self.nprobes = nprobes;
+        self
+    }
+
    /// Serve repeated reads of write-once table files from `cache`.
    pub fn with_read_cache(mut self, cache: ReadCache) -> Self {
        self.read_cache = Some(Arc::new(cache));
@@ -669,7 +677,7 @@ impl VectorService for VectorServer {
            },
            postfilter: req.postfilter,
            distance,
-            nprobes: None,
+            nprobes: self.nprobes,
            index_segments: None,
            fragments: None,
        };