vector: index-segment restriction + partition helper (P3 seam) (#12124)

* vector: allow restricting ANN search to index segments

Query gains an optional index_segments list; when set and the index
path is used, the scanner restricts the vector index search to those
segments via with_index_segments. This is the seam the Enterprise
planner uses to fan one query out across per-segment workers; OSS
keeps passing None.

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

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

* vector: enumerate and partition index segments for fanout

segments::index_segments lists the physical segment UUIDs covering the
vector field at the pinned version; partition() assigns them
round-robin to workers. The parity test proves two workers scanning
disjoint shares merge back to the full-dataset top-K — the P3 exit
gate, in-process.

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

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

* vector: let a query carry fragment ids for the unindexed sweep

Lance only applies the index-segment restriction to indexed rows —
unindexed fragments outside the selected segments are skipped. For a
partitioned query to stay correct, the planner needs to hand the
unindexed sweep to a worker too, so Query gains a fragments list that
the scanner resolves to manifest fragments via with_fragments. A task
carrying both searches its segments via the index and its fragments
with flat KNN in the same scan.

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

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

* vector: empty scope returns no hits; cover the sweep in tests

Some([]) fragments or index_segments fell through to an unrestricted
scan — a worker assigned nothing would have searched everything.
empty_scope short-circuits those queries, and the segment restriction
only calls with_index_segments on a non-empty list.

segments::segment_fragments also returns each segment's live fragment
bitmap, and the parity test now leaves a fragment unindexed so the
worker carrying it exercises the flat sweep — the contract the
Enterprise planner relies on.

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

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

* vector: scope the scan to listed ∪ segment-covered fragments

Lance intersects with_fragments with the selected index segments, so a
worker given its segment share plus the unindexed sweep lost the share
entirely — only the listed fragments were scanned. Expand the fragment
list with the live coverage of the selected segments. The list is now
also applied when it resolves empty: stale ids or an empty assignment
must scan nothing, not fall through to an unrestricted scan. The merge
test now leaves a fragment unindexed so the combined path is exercised.

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

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

* vector: partition assigns nothing when there are no workers

n.max(1) returned every segment in one phantom partition for a
nonexistent worker, silently dropping coverage.

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

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

---------

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

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

@@ -130,6 +130,8 @@ mod tests {
                 postfilter: false,
                 distance: DistanceType::L2,
                 nprobes: Some(4),
+                index_segments: None,
+                fragments: None,
             },
         )
         .await
@@ -156,6 +158,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: Some(4),
+            index_segments: None,
+            fragments: None,
         };
         let filtered = ann(&dataset, &table, &query).await.unwrap();
         let exact = flat(&dataset, &table, &query).await.unwrap();
@@ -191,6 +195,8 @@ mod tests {
                 postfilter: false,
                 distance: DistanceType::L2,
                 nprobes: None,
+                index_segments: None,
+                fragments: None,
             },
         )
         .await

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

@@ -8,6 +8,7 @@
 pub mod coverage;
 pub mod maintenance;
 pub mod query;
+pub mod segments;
 pub mod table;
 #[cfg(test)]
 mod testing;

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

@@ -233,6 +233,8 @@ async fn main() -> Result<()> {
                 postfilter,
                 distance,
                 nprobes,
+                index_segments: None,
+                fragments: None,
             };
             let hits = if flat {
                 query::flat(&dataset, &vector_table, &query).await?

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

@@ -172,6 +172,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         let mut ids: Vec<u64> = query::flat(&pinned, &resolved, &q)
             .await
@@ -414,6 +416,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         assert!(query::flat(&dataset, &resolved, &q).await.is_err());
     }

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

@@ -7,6 +7,7 @@ use futures::TryStreamExt;
 use lance::Dataset;
 use lance::dataset::scanner::Scanner;
 use lance_linalg::distance::DistanceType;
+use lance_table::format::Fragment;
 
 use crate::table::VectorTable;
 
@@ -27,6 +28,17 @@ pub struct Query {
     pub postfilter: bool,
     pub distance: DistanceType,
     pub nprobes: Option<usize>,
+    /// Restrict the indexed scan to these physical index segments (the
+    /// Enterprise planner fans a query out across them). Without `fragments`,
+    /// unindexed fragments outside the segments are not searched — the caller
+    /// owns disjoint coverage of them. When `fragments` is also set the scan
+    /// covers both, so a worker can take its segment share plus an
+    /// unindexed-fragment sweep in one request.
+    pub index_segments: Option<Vec<uuid::Uuid>>,
+    /// Restrict the scan to these fragment ids — used by the planner to give
+    /// a worker the unindexed-fragment sweep. Listed ids that no longer exist
+    /// scan nothing.
+    pub fragments: Option<Vec<u64>>,
 }
 
 /// One row of a top-K answer.
@@ -41,14 +53,14 @@ pub struct Hit {
 
 /// Exact top-K over every fragment, ignoring any index.
 pub async fn flat(dataset: &Dataset, table: &VectorTable, query: &Query) -> Result<Vec<Hit>> {
-    run(scanner(dataset, table, query, false)?, table).await
+    run(scanner(dataset, table, query, false).await?, table).await
 }
 
 /// Top-K with the ANN index where it covers the fragment and a flat pass over
 /// the rest. Whether the index is used is Lance's decision per fragment, so
 /// rows written after the last index build are always candidates.
 pub async fn ann(dataset: &Dataset, table: &VectorTable, query: &Query) -> Result<Vec<Hit>> {
-    run(scanner(dataset, table, query, true)?, table).await
+    run(scanner(dataset, table, query, true).await?, table).await
 }
 
 /// Fetch rows by key.
@@ -80,7 +92,7 @@ pub async fn get(
         .context("scan by keys")
 }
 
-fn scanner(
+async fn scanner(
     dataset: &Dataset,
     table: &VectorTable,
     query: &Query,
@@ -107,6 +119,34 @@ fn scanner(
     if let Some(nprobes) = query.nprobes {
         scanner.nprobes(nprobes);
     }
+    let segments = query.index_segments.as_deref().unwrap_or(&[]);
+    if use_index && !segments.is_empty() {
+        scanner.with_index_segments(segments.to_vec())?;
+    }
+    // Lance intersects the fragment list with the selected segments, so a
+    // scoped scan must name the segment-covered fragments itself: the universe
+    // is the listed ids plus the live coverage of the selected segments. An
+    // empty universe (or an empty-segment, fragment-less assignment) scopes to
+    // nothing — `with_fragments` with an empty list scans nothing.
+    let mut wanted: Option<Vec<u64>> = query.fragments.clone();
+    if let Some(ids) = &mut wanted {
+        if use_index && !segments.is_empty() {
+            ids.extend(
+                crate::segments::covered_fragments(dataset, &table.vector_column, segments).await?,
+            );
+        }
+    } else if use_index && query.index_segments.as_ref().is_some_and(|s| s.is_empty()) {
+        wanted = Some(Vec::new());
+    }
+    if let Some(ids) = wanted {
+        let frags: Vec<Fragment> = dataset
+            .fragments()
+            .iter()
+            .filter(|f| ids.contains(&f.id))
+            .cloned()
+            .collect();
+        scanner.with_fragments(frags);
+    }
     let mut columns = Vec::new();
     if let Some(key) = &table.key_column {
         columns.push(key.clone());
@@ -195,6 +235,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         let hits = flat(&dataset, &table, &query).await.unwrap();
         assert_eq!(hits.len(), 2);
@@ -234,6 +276,8 @@ mod tests {
                 postfilter: false,
                 distance: DistanceType::L2,
                 nprobes: None,
+                index_segments: None,
+                fragments: None,
             },
         )
         .await
@@ -311,6 +355,8 @@ mod tests {
                     postfilter: false,
                     distance: DistanceType::L2,
                     nprobes: None,
+                    index_segments: None,
+                    fragments: None,
                 },
             )
             .await
@@ -399,6 +445,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         assert!(flat(&dataset, &table, &query).await.is_err());
     }
@@ -420,6 +468,8 @@ mod tests {
             postfilter: false,
             distance: DistanceType::L2,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         let indexed = ann(&dataset, &table, &query).await.unwrap();
         let exact = flat(&dataset, &table, &query).await.unwrap();

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

@@ -0,0 +1,227 @@
+//! Physical index segments and the partition a fanout planner assigns to each
+//! worker. OSS runs the same restriction in-process; Enterprise ships the
+//! segment ids to remote workers — the engine does not change.
+
+use anyhow::Result;
+use lance::Dataset;
+use lance::index::DatasetIndexExt;
+use uuid::Uuid;
+
+/// UUIDs of every physical index segment covering `field` at this version.
+pub async fn index_segments(dataset: &Dataset, field: &str) -> Result<Vec<Uuid>> {
+    Ok(segment_fragments(dataset, field)
+        .await?
+        .iter()
+        .map(|(id, _)| *id)
+        .collect())
+}
+
+/// Per index segment on `field`: its uuid and the live fragments it covers.
+/// Segments with unknown coverage (no bitmap) are skipped — callers needing
+/// strict coverage must treat that as a degenerate plan.
+pub async fn segment_fragments(
+    dataset: &Dataset,
+    field: &str,
+) -> Result<Vec<(Uuid, roaring::RoaringBitmap)>> {
+    let field = dataset
+        .schema()
+        .field(field)
+        .ok_or_else(|| anyhow::anyhow!("no such field {field}"))?;
+    let live: roaring::RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect();
+    Ok(dataset
+        .load_indices()
+        .await?
+        .iter()
+        .filter(|i| i.fields == [field.id])
+        .filter_map(|i| i.fragment_bitmap.clone().map(|b| (i.uuid, b & &live)))
+        .collect())
+}
+
+/// Live fragment ids covered by `segments` on `field` — the union a scoped
+/// fragment list needs so the selected segments are not filtered out of the
+/// scan.
+pub async fn covered_fragments(
+    dataset: &Dataset,
+    field: &str,
+    segments: &[Uuid],
+) -> Result<Vec<u64>> {
+    Ok(segment_fragments(dataset, field)
+        .await?
+        .into_iter()
+        .filter(|(id, _)| segments.contains(id))
+        .flat_map(|(_, b)| b.iter().map(u64::from).collect::<Vec<_>>())
+        .collect())
+}
+
+/// Round-robin `segments` across `n` workers. Index order is unstable, so
+/// workers see disjoint, deterministic-in-name partitions rather than
+/// physically adjacent ones. Zero workers means zero partitions — nothing is
+/// silently assigned.
+pub fn partition(segments: &[Uuid], n: usize) -> Vec<Vec<Uuid>> {
+    let mut out = vec![Vec::new(); n];
+    if n == 0 {
+        return out;
+    }
+    let mut ids = segments.to_vec();
+    ids.sort();
+    for (i, id) in ids.into_iter().enumerate() {
+        out[i % n].push(id);
+    }
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{Query, VectorTable, query, testing};
+    use lance::index::vector::VectorIndexParams;
+    use lance_index::IndexType;
+    use lance_index::optimize::OptimizeOptions;
+    use lance_linalg::distance::DistanceType;
+    use tempfile::TempDir;
+
+    /// Two workers each scan their share of the index segments; merged, the
+    /// hits equal a single full-dataset search.
+    #[tokio::test]
+    async fn partitioned_segments_merge_to_full_result() {
+        let dir = TempDir::new().unwrap();
+        let mut dataset = testing::write_dataset(
+            &dir,
+            &testing::vectors(8, 4),
+            &["a", "b", "c", "d", "e", "f", "g", "h"],
+        )
+        .await;
+        let params = VectorIndexParams::ivf_flat(4, DistanceType::L2);
+        dataset
+            .create_index(
+                &["vector"],
+                IndexType::Vector,
+                Some("idx".to_string()),
+                &params,
+                false,
+            )
+            .await
+            .unwrap();
+        // Second fragment + delta segment so the index has two parts, then a
+        // third fragment left unindexed for the sweep half of the plan.
+        testing::append_dataset(&mut dataset, &testing::vectors(4, 4), &["i", "j", "k", "l"]).await;
+        dataset
+            .optimize_indices(&OptimizeOptions::append())
+            .await
+            .unwrap();
+        testing::append_dataset(&mut dataset, &testing::vectors(2, 4), &["m", "n"]).await;
+
+        let resolved = VectorTable::resolve(&dataset, None, None).unwrap();
+        let base = Query {
+            vector: vec![0.5; 4],
+            k: 5,
+            filter: None,
+            postfilter: false,
+            distance: DistanceType::L2,
+            nprobes: None,
+            index_segments: None,
+            fragments: None,
+        };
+        let want: Vec<(u64, u32)> = query::ann(&dataset, &resolved, &base)
+            .await
+            .unwrap()
+            .iter()
+            .map(|h| (h.row_id, h.distance.to_bits()))
+            .collect();
+
+        let segments_bitmap = segment_fragments(&dataset, "vector").await.unwrap();
+        let segments: Vec<Uuid> = segments_bitmap.iter().map(|(id, _)| *id).collect();
+        assert!(
+            segments.len() >= 2,
+            "expected >=2 segments, got {segments:?}"
+        );
+        // The unindexed fragment belongs to nobody's segment; worker 0 sweeps
+        // it next to its own share — the same contract plan_search implements
+        // on the enterprise side.
+        let unindexed: Vec<u64> = {
+            let all: std::collections::BTreeSet<u64> =
+                dataset.fragments().iter().map(|f| f.id).collect();
+            let mut indexed = std::collections::BTreeSet::new();
+            for (_, b) in &segments_bitmap {
+                indexed.extend(b.iter().map(|i| i as u64));
+            }
+            all.difference(&indexed).copied().collect()
+        };
+        assert!(!unindexed.is_empty(), "expected an unindexed fragment");
+        let mut merged: Vec<(f32, u64)> = Vec::new();
+        for (w, share) in partition(&segments, 2).into_iter().enumerate() {
+            let mut worker_q = base.clone();
+            worker_q.index_segments = Some(share);
+            if w == 0 {
+                worker_q.fragments = Some(unindexed.clone());
+            }
+            merged.extend(
+                query::ann(&dataset, &resolved, &worker_q)
+                    .await
+                    .unwrap()
+                    .iter()
+                    .map(|h| (h.distance, h.row_id)),
+            );
+        }
+        merged.sort_by(|a, b| {
+            a.0.partial_cmp(&b.0)
+                .unwrap_or(std::cmp::Ordering::Equal)
+                .then(a.1.cmp(&b.1))
+        });
+        let got: Vec<(u64, u32)> = merged
+            .iter()
+            .take(base.k)
+            .map(|(d, id)| (*id, d.to_bits()))
+            .collect();
+        assert_eq!(got, want);
+    }
+
+    /// A scope of fragment ids that no longer exist, an empty list, or an
+    /// empty-segment assignment must return no hits — never fall through to
+    /// an unrestricted scan.
+    #[tokio::test]
+    async fn empty_scope_scans_nothing() {
+        let dir = TempDir::new().unwrap();
+        let dataset =
+            testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await;
+        let resolved = VectorTable::resolve(&dataset, None, None).unwrap();
+        let mut q = Query {
+            vector: vec![0.5; 4],
+            k: 4,
+            filter: None,
+            postfilter: false,
+            distance: DistanceType::L2,
+            nprobes: None,
+            index_segments: None,
+            fragments: Some(vec![u64::MAX]),
+        };
+        assert!(
+            query::ann(&dataset, &resolved, &q)
+                .await
+                .unwrap()
+                .is_empty()
+        );
+        assert!(
+            query::flat(&dataset, &resolved, &q)
+                .await
+                .unwrap()
+                .is_empty()
+        );
+        q.fragments = Some(Vec::new());
+        assert!(
+            query::ann(&dataset, &resolved, &q)
+                .await
+                .unwrap()
+                .is_empty()
+        );
+        q.fragments = None;
+        q.index_segments = Some(Vec::new());
+        assert!(
+            query::ann(&dataset, &resolved, &q)
+                .await
+                .unwrap()
+                .is_empty()
+        );
+        assert!(partition(&[Uuid::nil()], 0).is_empty());
+    }
+}

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

@@ -42,6 +42,8 @@ fn query_for(vector: Vec<f32>, k: usize) -> Query {
         postfilter: false,
         distance: DistanceType::L2,
         nprobes: Some(4),
+        index_segments: None,
+        fragments: None,
     }
 }
 

seaweed-vector/crates/server/src/service.rs

@@ -76,17 +76,16 @@ fn schema_for(
     );
     // The arrow.json extension marks the Utf8 column as JSON: lance stores it
     // as jsonb, which is what the json_get_* filter functions take.
-    let mut metadata_meta = HashMap::from_iter([(
-        "ARROW:extension:name".to_string(),
-        "arrow.json".to_string(),
-    )]);
+    let mut metadata_meta =
+        HashMap::from_iter([("ARROW:extension:name".to_string(), "arrow.json".to_string())]);
     if !non_filterable.is_empty() {
         metadata_meta.insert(
             "non_filterable_metadata_keys".to_string(),
             serde_json::to_string(non_filterable).unwrap_or_default(),
         );
     }
-    let metadata_field = Field::new(METADATA_COLUMN, DataType::Utf8, true).with_metadata(metadata_meta);
+    let metadata_field =
+        Field::new(METADATA_COLUMN, DataType::Utf8, true).with_metadata(metadata_meta);
     Arc::new(Schema::new(vec![
         Field::new(KEY_COLUMN, DataType::Utf8, false),
         vector_field,
@@ -122,12 +121,8 @@ fn batch_of(vectors: &[pb::Vector], dimension: usize, schema: Arc<Schema>) -> Re
     else {
         bail!("vector column is not a fixed-size list");
     };
-    let vectors_array = FixedSizeListArray::try_new(
-        item_field.clone(),
-        dimension as i32,
-        values,
-        None,
-    )?;
+    let vectors_array =
+        FixedSizeListArray::try_new(item_field.clone(), dimension as i32, values, None)?;
     let metas = vectors
         .iter()
         .map(|v| {
@@ -523,6 +518,8 @@ impl VectorService for VectorServer {
             postfilter: req.postfilter,
             distance,
             nprobes: None,
+            index_segments: None,
+            fragments: None,
         };
         let hits = query::ann(&dataset, &resolved, &query)
             .await