vector: ANN coverage introspection + prefilter (P0 phase 2) (#12097)

* vector: share the test dataset builder

The coverage tests need the same tempdir Lance datasets the query tests
build; move them into a cfg(test) helper module.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: prefilter the metadata filter

prefilter=false post-filters the top-K rows, which can return fewer than
k results and never sees unindexed rows beyond the nearest cut. Expose
the choice so the S3 Vectors layer can honour AWS filter semantics.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: index coverage introspection

coverage() unions each index's fragment bitmaps against the manifest's
fragment list to report which fragments a vector index covers, which it
has not reached, and how many rows the flat pass still has to scan.
Tests pin the behaviour the spike is for: rows appended after an index
build stay in the top-K, and a version-pinned dataset never sees the
next commit.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: default filtered queries to prefilter

postfilter=false is now the normal path: the metadata filter cuts the
eligible rows before ranking, so a query returns k hits whenever k
eligible rows exist. Lance's post-filter remains reachable through
Query::postfilter for comparison runs.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: intersect index bitmaps with the live manifest

A bitmap records fragments the index covered when built; a fragment the
manifest dropped afterwards stayed in the reported coverage.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: keep bitmap-less indices in coverage

An index whose metadata has no fragment bitmap was dropped from the
report entirely, indistinguishable from no index at all. Count its
segments and mark its coverage unknown instead.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: count live rows in unindexed coverage

physical_rows counts deleted rows and can be absent in the manifest;
sum FileFragment::count_rows for the unindexed fragments instead.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* vector: error on a null key in a selected key column

* vector: compare filtered ANN and flat hits as sets

* vector: coverage reports the checked-out version

* vector: project _distance and _rowid explicitly

---------

Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: chrislusf <chris@seaweedfs.com>

seaweed-vector/crates/core/src/coverage.rs

@@ -0,0 +1,205 @@
+//! Which fragments the vector index covers, and which it has not reached yet.
+
+use std::collections::BTreeMap;
+
+use anyhow::{Context, Result};
+use lance::Dataset;
+use lance::index::DatasetIndexExt;
+use roaring::RoaringBitmap;
+
+/// How much of the dataset one named index covers at this version.
+#[derive(Debug)]
+pub struct IndexCoverage {
+    pub index_name: String,
+    /// Physical index segments: the IndexMetadata deltas sharing the name.
+    pub segments: usize,
+    /// False when any delta of this index lacks a fragment bitmap: the index
+    /// exists but which fragments it covers is unknown.
+    pub coverage_known: bool,
+    pub indexed_fragments: Vec<u64>,
+    pub unindexed_fragments: Vec<u64>,
+    pub unindexed_rows: usize,
+}
+
+/// Coverage for every index built on `vector_column` at this version.
+pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result<Vec<IndexCoverage>> {
+    let field = dataset
+        .schema()
+        .field(vector_column)
+        .with_context(|| format!("no such column {vector_column}"))?;
+    let mut by_name: BTreeMap<String, Vec<Option<RoaringBitmap>>> = BTreeMap::new();
+    for index in dataset.load_indices().await?.iter() {
+        if index.fields != [field.id] {
+            continue;
+        }
+        by_name
+            .entry(index.name.clone())
+            .or_default()
+            .push(index.fragment_bitmap.clone());
+    }
+    let all: RoaringBitmap = dataset
+        .fragments()
+        .iter()
+        .map(|f| f.id as u32)
+        .collect();
+    let mut out = Vec::new();
+    for (name, bitmaps) in by_name {
+        let coverage_known = bitmaps.iter().all(|b| b.is_some());
+        let mut indexed = RoaringBitmap::new();
+        for bitmap in bitmaps.iter().flatten() {
+            indexed |= bitmap;
+        }
+        // The bitmap records what the index covered when built; fragments
+        // the manifest has since dropped are not part of this version.
+        let indexed = indexed & &all;
+        let unindexed = &all - &indexed;
+        let mut unindexed_rows = 0;
+        for fragment in dataset.get_fragments() {
+            // physical_rows counts deleted rows too, and may be absent in the
+            // manifest; count live rows instead.
+            if unindexed.contains(fragment.id() as u32) {
+                unindexed_rows += fragment.count_rows(None).await?;
+            }
+        }
+        out.push(IndexCoverage {
+            index_name: name,
+            segments: bitmaps.len(),
+            coverage_known,
+            indexed_fragments: indexed.iter().map(|id| id as u64).collect(),
+            unindexed_rows,
+            unindexed_fragments: unindexed.iter().map(|id| id as u64).collect(),
+        });
+    }
+    Ok(out)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::query::{Query, ann, flat};
+    use crate::table::VectorTable;
+    use crate::testing::{DIM, append_dataset, vectors, write_dataset};
+    use lance::index::vector::VectorIndexParams;
+    use lance_index::IndexType;
+    use lance_linalg::distance::DistanceType;
+    use tempfile::TempDir;
+
+    async fn indexed_dataset(dir: &TempDir, tags: &[&str]) -> Dataset {
+        let mut dataset = write_dataset(dir, &vectors(tags.len(), DIM), tags).await;
+        let params = VectorIndexParams::ivf_flat(4, DistanceType::L2);
+        dataset
+            .create_index(
+                &["vector"],
+                IndexType::Vector,
+                Some("vidx".to_string()),
+                &params,
+                true,
+            )
+            .await
+            .unwrap();
+        dataset.checkout_latest().await.unwrap();
+        dataset
+    }
+
+    #[tokio::test]
+    async fn coverage_reports_the_unindexed_fragment() {
+        let dir = TempDir::new().unwrap();
+        let mut dataset = indexed_dataset(&dir, &["a"; 64]).await;
+        append_dataset(&mut dataset, &vectors(4, DIM), &["b"; 4]).await;
+
+        let coverages = coverage(&dataset, "vector").await.unwrap();
+        assert_eq!(coverages.len(), 1);
+        let cov = &coverages[0];
+        assert_eq!(cov.index_name, "vidx");
+        assert_eq!(cov.segments, 1);
+        assert_eq!(cov.indexed_fragments, vec![0]);
+        assert_eq!(cov.unindexed_fragments.len(), 1);
+        assert_eq!(cov.unindexed_rows, 4);
+    }
+
+    #[tokio::test]
+    async fn ann_searches_unindexed_fragments() {
+        let dir = TempDir::new().unwrap();
+        let mut dataset = indexed_dataset(&dir, &["a"; 64]).await;
+        append_dataset(&mut dataset, &[vec![9.0; DIM]], &["b"]).await;
+        let table = VectorTable::resolve(&dataset, None, None).unwrap();
+
+        let hits = ann(
+            &dataset,
+            &table,
+            &Query {
+                vector: vec![9.0; DIM],
+                k: 1,
+                filter: None,
+                postfilter: false,
+                distance: DistanceType::L2,
+                nprobes: Some(4),
+            },
+        )
+        .await
+        .unwrap();
+        assert_eq!(hits.len(), 1);
+        assert_eq!(hits[0].key, "k64");
+        assert!(hits[0].distance < 1e-6);
+    }
+
+    #[tokio::test]
+    async fn prefiltered_ann_matches_filtered_flat() {
+        let dir = TempDir::new().unwrap();
+        let tags: Vec<&str> = (0..64)
+            .map(|i| if i % 2 == 0 { "a" } else { "b" })
+            .collect();
+        let mut dataset = indexed_dataset(&dir, &tags).await;
+        append_dataset(&mut dataset, &vectors(4, DIM), &["a"; 4]).await;
+        let table = VectorTable::resolve(&dataset, None, None).unwrap();
+
+        let query = Query {
+            vector: vectors(1, DIM).remove(0),
+            k: 10,
+            filter: Some("tag = 'a'".to_string()),
+            postfilter: false,
+            distance: DistanceType::L2,
+            nprobes: Some(4),
+        };
+        let filtered = ann(&dataset, &table, &query).await.unwrap();
+        let exact = flat(&dataset, &table, &query).await.unwrap();
+        // equal-distance hits have no specified order between the two scans
+        let mut filtered_keys = filtered.iter().map(|h| &h.key).collect::<Vec<_>>();
+        let mut exact_keys = exact.iter().map(|h| &h.key).collect::<Vec<_>>();
+        filtered_keys.sort();
+        exact_keys.sort();
+        assert_eq!(filtered_keys, exact_keys);
+        assert!(filtered.len() <= 10);
+        assert!(filtered.iter().all(|h| h.key != "k1"));
+    }
+
+    #[tokio::test]
+    async fn pinned_version_ignores_newer_commits() {
+        let dir = TempDir::new().unwrap();
+        let mut dataset = write_dataset(&dir, &vectors(4, DIM), &["a"; 4]).await;
+        let pinned = dataset
+            .checkout_version(dataset.version().version)
+            .await
+            .unwrap();
+        append_dataset(&mut dataset, &vectors(4, DIM), &["b"; 4]).await;
+
+        assert_eq!(pinned.count_rows(None).await.unwrap(), 4);
+        let table = VectorTable::resolve(&pinned, None, None).unwrap();
+        let hits = flat(
+            &pinned,
+            &table,
+            &Query {
+                vector: vectors(1, DIM).remove(0),
+                k: 10,
+                filter: None,
+                postfilter: false,
+                distance: DistanceType::L2,
+                nprobes: None,
+            },
+        )
+        .await
+        .unwrap();
+        assert_eq!(hits.len(), 4);
+        assert!(hits.iter().all(|h| h.row_id < 4));
+    }
+}

seaweed-vector/crates/core/src/lib.rs

@@ -5,8 +5,12 @@
 //! the whole dataset: index-accelerated where an ANN index exists, exact where
 //! it does not, always covering the fragments no index has reached yet.
 
+pub mod coverage;
 pub mod query;
 pub mod table;
+#[cfg(test)]
+mod testing;
 
+pub use coverage::IndexCoverage;
 pub use query::{Hit, Query};
 pub use table::VectorTable;

seaweed-vector/crates/core/src/query.rs

@@ -20,6 +20,11 @@ pub struct Query {
     pub k: usize,
     /// A Lance SQL filter over the metadata columns.
     pub filter: Option<String>,
+    /// Filter the top-K rows after ranking instead of the eligible rows
+    /// before it. Lance's default post-filters and can return fewer than k
+    /// results while eligible neighbours go unreported; keep it off for
+    /// S3 Vectors filter semantics.
+    pub postfilter: bool,
     pub distance: DistanceType,
     pub nprobes: Option<usize>,
 }
@@ -97,6 +102,7 @@ fn scanner(
         .with_row_id();
     if let Some(filter) = &query.filter {
         scanner.filter(filter)?;
+        scanner.prefilter(!query.postfilter);
     }
     if let Some(nprobes) = query.nprobes {
         scanner.nprobes(nprobes);
@@ -105,6 +111,8 @@ fn scanner(
     if let Some(key) = &table.key_column {
         columns.push(key.clone());
     }
+    columns.push(DISTANCE_COLUMN.to_string());
+    columns.push(ROW_ID_COLUMN.to_string());
     scanner.project(&columns)?;
     Ok(scanner)
 }
@@ -137,11 +145,16 @@ fn hits(batch: &RecordBatch, table: &VectorTable) -> Result<Vec<Hit>> {
     }
     (0..batch.num_rows())
         .map(|i| {
+            let key = match keys {
+                Some(keys) if !keys.is_null(i) => keys.value(i).to_string(),
+                Some(_) => bail!(
+                    "key column {} has a null at row {i}",
+                    table.key_column.as_deref().unwrap_or_default()
+                ),
+                None => row_ids.value(i).to_string(),
+            };
             Ok(Hit {
-                key: match keys {
-                    Some(keys) if !keys.is_null(i) => keys.value(i).to_string(),
-                    _ => row_ids.value(i).to_string(),
-                },
+                key,
                 distance: distances.value(i),
                 row_id: row_ids.value(i),
             })
@@ -153,53 +166,9 @@ fn hits(batch: &RecordBatch, table: &VectorTable) -> Result<Vec<Hit>> {
 mod tests {
     use super::*;
     use crate::table::VectorTable;
-    use arrow_array::{FixedSizeListArray, RecordBatchIterator, StringArray, types::Float32Type};
-    use arrow_schema::{DataType, Field, Schema};
-    use std::sync::Arc;
+    use crate::testing::{DIM, write_dataset};
     use tempfile::TempDir;
 
-    const DIM: usize = 4;
-
-    fn vector_column(vectors: &[Vec<f32>]) -> FixedSizeListArray {
-        FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
-            vectors
-                .iter()
-                .map(|v| Some(v.iter().copied().map(Some).collect::<Vec<_>>())),
-            DIM as i32,
-        )
-    }
-
-    async fn write_dataset(dir: &TempDir, vectors: &[Vec<f32>], tags: &[&str]) -> Dataset {
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("key", DataType::Utf8, false),
-            Field::new(
-                "vector",
-                DataType::FixedSizeList(
-                    Arc::new(Field::new("item", DataType::Float32, true)),
-                    DIM as i32,
-                ),
-                false,
-            ),
-            Field::new("tag", DataType::Utf8, true),
-        ]));
-        let keys: Vec<String> = (0..vectors.len()).map(|i| format!("k{i}")).collect();
-        let batch = RecordBatch::try_new(
-            schema.clone(),
-            vec![
-                Arc::new(StringArray::from(keys)),
-                Arc::new(vector_column(vectors)),
-                Arc::new(StringArray::from(
-                    tags.iter().map(|t| Some(t.to_string())).collect::<Vec<_>>(),
-                )),
-            ],
-        )
-        .unwrap();
-        let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
-        Dataset::write(reader, dir.path().to_str().unwrap(), None)
-            .await
-            .unwrap()
-    }
-
     fn l2(a: &[f32], b: &[f32]) -> f32 {
         a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum()
     }
@@ -220,6 +189,7 @@ mod tests {
             vector: vec![1.0, 0.0, 0.0, 0.0],
             k: 2,
             filter: None,
+            postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
         };
@@ -258,6 +228,7 @@ mod tests {
                 vector: vec![1.0, 0.0, 0.0, 0.0],
                 k: 3,
                 filter: Some("tag = 'a'".to_string()),
+                postfilter: false,
                 distance: DistanceType::L2,
                 nprobes: None,
             },
@@ -293,6 +264,7 @@ mod tests {
             vector: vec![1.0; DIM + 1],
             k: 1,
             filter: None,
+            postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
         };
@@ -313,6 +285,7 @@ mod tests {
             vector: vec![1.0, 0.0, 0.0, 0.0],
             k: 2,
             filter: None,
+            postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
         };

seaweed-vector/crates/core/src/testing.rs

@@ -0,0 +1,93 @@
+//! Small Lance datasets for the query and coverage tests.
+
+use std::sync::Arc;
+
+use arrow_array::{
+    FixedSizeListArray, RecordBatch, RecordBatchIterator, StringArray, types::Float32Type,
+};
+use arrow_schema::{DataType, Field, Schema};
+use lance::Dataset;
+use lance::dataset::WriteParams;
+use tempfile::TempDir;
+
+pub const DIM: usize = 4;
+
+fn vector_column(vectors: &[Vec<f32>], dim: usize) -> FixedSizeListArray {
+    FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
+        vectors
+            .iter()
+            .map(|v| Some(v.iter().copied().map(Some).collect::<Vec<_>>())),
+        dim as i32,
+    )
+}
+
+pub fn schema(dim: usize) -> Arc<Schema> {
+    Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Utf8, false),
+        Field::new(
+            "vector",
+            DataType::FixedSizeList(
+                Arc::new(Field::new("item", DataType::Float32, true)),
+                dim as i32,
+            ),
+            false,
+        ),
+        Field::new("tag", DataType::Utf8, true),
+    ]))
+}
+
+pub fn batch(offset: usize, vectors: &[Vec<f32>], tags: &[&str], dim: usize) -> RecordBatch {
+    let keys: Vec<String> = (0..vectors.len())
+        .map(|i| format!("k{}", offset + i))
+        .collect();
+    RecordBatch::try_new(
+        schema(dim),
+        vec![
+            Arc::new(StringArray::from(keys)),
+            Arc::new(vector_column(vectors, dim)),
+            Arc::new(StringArray::from(
+                tags.iter().map(|t| Some(t.to_string())).collect::<Vec<_>>(),
+            )),
+        ],
+    )
+    .unwrap()
+}
+
+pub async fn write_dataset(dir: &TempDir, vectors: &[Vec<f32>], tags: &[&str]) -> Dataset {
+    write_dataset_at(dir.path().to_str().unwrap(), vectors, tags).await
+}
+
+pub async fn write_dataset_at(uri: &str, vectors: &[Vec<f32>], tags: &[&str]) -> Dataset {
+    let dim = vectors[0].len();
+    let reader = RecordBatchIterator::new(
+        vec![Ok(batch(0, vectors, tags, dim))].into_iter(),
+        schema(dim),
+    );
+    Dataset::write(reader, uri, None).await.unwrap()
+}
+
+/// Deterministic non-trivial vectors for index coverage tests.
+pub fn vectors(rows: usize, dim: usize) -> Vec<Vec<f32>> {
+    (0..rows)
+        .map(|i| {
+            (0..dim)
+                .map(|j| ((i * 31 + j * 17) % 97) as f32 / 97.0)
+                .collect()
+        })
+        .collect()
+}
+
+/// Commit a new version whose fragment is not yet covered by the index the
+/// dataset had before.
+pub async fn append_dataset(dataset: &mut Dataset, vectors: &[Vec<f32>], tags: &[&str]) {
+    let offset = dataset.count_rows(None).await.unwrap();
+    let dim = vectors[0].len();
+    let reader = RecordBatchIterator::new(
+        vec![Ok(batch(offset, vectors, tags, dim))].into_iter(),
+        schema(dim),
+    );
+    dataset
+        .append(reader, Some(WriteParams::default()))
+        .await
+        .unwrap();
+}