From af77406cf9513b41197d08210a3aa972debba19e Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Tue, 27 Oct 2026 02:22:22 -0700 Subject: [PATCH] vector: index-segment restriction + partition helper (P3 seam) (#12124) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 | 6 + seaweed-vector/crates/core/src/lib.rs | 1 + seaweed-vector/crates/core/src/main.rs | 2 + seaweed-vector/crates/core/src/maintenance.rs | 4 + seaweed-vector/crates/core/src/query.rs | 56 ++++- seaweed-vector/crates/core/src/segments.rs | 227 ++++++++++++++++++ seaweed-vector/crates/core/src/verify.rs | 2 + seaweed-vector/crates/server/src/service.rs | 19 +- 8 files changed, 303 insertions(+), 14 deletions(-) create mode 100644 seaweed-vector/crates/core/src/segments.rs diff --git a/seaweed-vector/crates/core/src/coverage.rs b/seaweed-vector/crates/core/src/coverage.rs index ad7d8e650a..7849215497 100644 --- a/seaweed-vector/crates/core/src/coverage.rs +++ b/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 diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs index f5ce66ad75..def0eb215f 100644 --- a/seaweed-vector/crates/core/src/lib.rs +++ b/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; diff --git a/seaweed-vector/crates/core/src/main.rs b/seaweed-vector/crates/core/src/main.rs index 48f42f9b77..2c3d745848 100644 --- a/seaweed-vector/crates/core/src/main.rs +++ b/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? diff --git a/seaweed-vector/crates/core/src/maintenance.rs b/seaweed-vector/crates/core/src/maintenance.rs index 04fc8029d6..b75b521330 100644 --- a/seaweed-vector/crates/core/src/maintenance.rs +++ b/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 = 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()); } diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs index 89c33f3f0f..63ffd72a9c 100644 --- a/seaweed-vector/crates/core/src/query.rs +++ b/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, + /// 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>, + /// 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>, } /// 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> { - 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> { - 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> = 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 = 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(); diff --git a/seaweed-vector/crates/core/src/segments.rs b/seaweed-vector/crates/core/src/segments.rs new file mode 100644 index 0000000000..e7767efafd --- /dev/null +++ b/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> { + 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> { + 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> { + Ok(segment_fragments(dataset, field) + .await? + .into_iter() + .filter(|(id, _)| segments.contains(id)) + .flat_map(|(_, b)| b.iter().map(u64::from).collect::>()) + .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> { + 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()), + ¶ms, + 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 = 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 = { + let all: std::collections::BTreeSet = + 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()); + } +} diff --git a/seaweed-vector/crates/core/src/verify.rs b/seaweed-vector/crates/core/src/verify.rs index 2ba0b91b8d..33eaea498d 100644 --- a/seaweed-vector/crates/core/src/verify.rs +++ b/seaweed-vector/crates/core/src/verify.rs @@ -42,6 +42,8 @@ fn query_for(vector: Vec, k: usize) -> Query { postfilter: false, distance: DistanceType::L2, nprobes: Some(4), + index_segments: None, + fragments: None, } } diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index 376dea4edd..f94dbc9415 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/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) -> 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