From 7388d0d4fc7e5161cbc33f8bf61633efed062669 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Mon, 26 Oct 2026 10:13:05 -0700 Subject: [PATCH] vector: seaweed-vector workspace + flat/indexed top-K engine (P0 phase 1) (#12096) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * vector: scaffold the seaweed-vector workspace Standalone cargo workspace for the vector query engine, mirroring the seaweed-worker layout: own lockfile, same lint policy, path dependencies across the two trees since there is no common root manifest. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: resolve a table's vector layout VectorTable::resolve finds the vector column (explicit name, else the column named vector, else the sole fixed-size float list), the optional key column, and the metadata columns a query can project. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: flat and indexed top-k, key lookup query::flat runs an exact scan (use_index(false)); query::ann lets Lance split the query between index segments and a flat pass over fragments no index covers yet, so rows written after the last index build stay eligible. Both honour a metadata filter and return the raw _distance in the query's metric. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: reject a non-utf8 or non-identifier key column An explicit --key-column naming a non-UTF8 field made hits() silently report row ids as keys, and get() interpolated the name into a SQL filter verbatim. Require Utf8 at resolve and in hits(), and require a plain identifier before building the IN predicate. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: fall back to row id on a null key keys.value() on a null slot fabricates an empty-string key; a row with no key reports its row id instead, matching the no-key-column case. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: quote the key column in the IN filter A Utf8 key column named e.g. document-key made the predicate parse as subtraction, and an identifier interpolated unquoted is injectable. Backtick-quote it and double its inner backticks; double quotes are string literals in the filter dialect, not identifiers. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: project only the columns a hit needs The knn projection pulled every metadata column — including other vector columns — although Hit returns none of them. Filters still evaluate off the scan path, so restrict the output to the key and the scoring columns. 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/Cargo.toml | 28 ++ seaweed-vector/crates/core/Cargo.toml | 34 +++ seaweed-vector/crates/core/src/lib.rs | 12 + seaweed-vector/crates/core/src/query.rs | 326 ++++++++++++++++++++++++ seaweed-vector/crates/core/src/table.rs | 178 +++++++++++++ 5 files changed, 578 insertions(+) create mode 100644 seaweed-vector/Cargo.toml create mode 100644 seaweed-vector/crates/core/Cargo.toml create mode 100644 seaweed-vector/crates/core/src/lib.rs create mode 100644 seaweed-vector/crates/core/src/query.rs create mode 100644 seaweed-vector/crates/core/src/table.rs diff --git a/seaweed-vector/Cargo.toml b/seaweed-vector/Cargo.toml new file mode 100644 index 0000000000..93fe4ff858 --- /dev/null +++ b/seaweed-vector/Cargo.toml @@ -0,0 +1,28 @@ +# Vector query engine for SeaweedFS Lance table buckets. +# +# `core` is the engine: open a table through the Lance namespace and run +# top-K over it. The crate layout follows vector_buckets.md; a `server` +# crate with the VectorService gRPC endpoint lands with the S3 Vectors +# milestone. +[workspace] +resolver = "2" +members = ["crates/core"] + +[workspace.package] +version = "0.1.0" +edition = "2024" +# The edition needs 1.85; the lance dependency tree needs more. Same floor as +# seaweed-worker, which this workspace path-depends on. +rust-version = "1.94.1" + +# No common root manifest links this tree to seaweed-worker or +# seaweed-common; keep the lint policy identical to theirs. +[workspace.lints.clippy] +needless_update = "allow" +undocumented_unsafe_blocks = "warn" + +[workspace.dependencies] +anyhow = "1" +tokio = { version = "1", features = ["full"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/seaweed-vector/crates/core/Cargo.toml b/seaweed-vector/crates/core/Cargo.toml new file mode 100644 index 0000000000..4241f4a136 --- /dev/null +++ b/seaweed-vector/crates/core/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "weed-vector-core" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +description = "Vector query engine over SeaweedFS Lance table buckets" + +[lib] +name = "weed_vector_core" + +[dependencies] +# The catalog client and dataset-opening helper the maintenance worker uses. +# A path dependency because the two trees are separate cargo workspaces with +# no common root manifest. +weed-lance-worker = { path = "../../../seaweed-worker/crates/lance" } +lance = { version = "10", default-features = false, features = ["aws"] } +lance-index = "10" +lance-linalg = "10" +lance-table = "10" +arrow-array = "58" +arrow-schema = "58" +roaring = "0.11" +anyhow.workspace = true +futures = "0.3" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio.workspace = true +tracing.workspace = true + +[dev-dependencies] +tempfile = "3" + +[lints] +workspace = true diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs new file mode 100644 index 0000000000..57c34a8d5e --- /dev/null +++ b/seaweed-vector/crates/core/src/lib.rs @@ -0,0 +1,12 @@ +//! Vector queries over the Lance tables the SeaweedFS namespace registers. +//! +//! A table bucket's rows carry a key column, a fixed-size float vector column, +//! and metadata columns. This crate resolves that layout and runs top-K over +//! 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 query; +pub mod table; + +pub use query::{Hit, Query}; +pub use table::VectorTable; diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs new file mode 100644 index 0000000000..49f8e207c9 --- /dev/null +++ b/seaweed-vector/crates/core/src/query.rs @@ -0,0 +1,326 @@ +//! Top-K queries over a [`VectorTable`]: exact flat scan, or the ANN index +//! with a flat pass over the fragments it does not cover yet. + +use anyhow::{Context, Result, bail}; +use arrow_array::{Array, Float32Array, RecordBatch, StringArray, UInt64Array}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::dataset::scanner::Scanner; +use lance_linalg::distance::DistanceType; + +use crate::table::VectorTable; + +const DISTANCE_COLUMN: &str = "_distance"; +const ROW_ID_COLUMN: &str = "_rowid"; + +/// One vector search request. +#[derive(Debug, Clone)] +pub struct Query { + pub vector: Vec, + pub k: usize, + /// A Lance SQL filter over the metadata columns. + pub filter: Option, + pub distance: DistanceType, + pub nprobes: Option, +} + +/// One row of a top-K answer. +#[derive(Debug, Clone)] +pub struct Hit { + /// The key column value, or the row id when the table has no key column. + pub key: String, + /// Raw Lance `_distance` in the query's metric. + pub distance: f32, + pub row_id: u64, +} + +/// 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 +} + +/// 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 +} + +/// Fetch rows by key. +pub async fn get( + dataset: &Dataset, + table: &VectorTable, + keys: &[String], +) -> Result> { + let Some(key_column) = &table.key_column else { + bail!("the table has no key column to look up"); + }; + if keys.is_empty() { + return Ok(Vec::new()); + } + let list = keys + .iter() + .map(|k| format!("'{}'", k.replace('\'', "''"))) + .collect::>() + .join(", "); + let mut scanner = dataset.scan(); + scanner + .filter(&format!("`{}` in ({list})", key_column.replace('`', "``")))? + .with_row_id(); + scanner + .try_into_stream() + .await? + .try_collect() + .await + .context("scan by keys") +} + +fn scanner( + dataset: &Dataset, + table: &VectorTable, + query: &Query, + use_index: bool, +) -> Result { + if query.vector.len() != table.dimension { + bail!( + "query dim({}) does not match the table's dim({})", + query.vector.len(), + table.dimension + ); + } + let vector = Float32Array::from(query.vector.clone()); + let mut scanner = dataset.scan(); + scanner + .nearest(&table.vector_column, &vector, query.k)? + .distance_metric(query.distance) + .use_index(use_index) + .with_row_id(); + if let Some(filter) = &query.filter { + scanner.filter(filter)?; + } + if let Some(nprobes) = query.nprobes { + scanner.nprobes(nprobes); + } + let mut columns = Vec::new(); + if let Some(key) = &table.key_column { + columns.push(key.clone()); + } + scanner.project(&columns)?; + Ok(scanner) +} + +async fn run(scanner: Scanner, table: &VectorTable) -> Result> { + let mut stream = scanner.try_into_stream().await?; + let mut out = Vec::new(); + while let Some(batch) = stream.try_next().await? { + out.extend(hits(&batch, table)?); + } + Ok(out) +} + +fn hits(batch: &RecordBatch, table: &VectorTable) -> Result> { + let distances = batch + .column_by_name(DISTANCE_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()) + .context("the scan did not return a _distance column")?; + let row_ids = batch + .column_by_name(ROW_ID_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()) + .context("the scan did not return a _rowid column")?; + let keys = table.key_column.as_ref().and_then(|name| { + batch + .column_by_name(name) + .and_then(|c| c.as_any().downcast_ref::()) + }); + if table.key_column.is_some() && keys.is_none() { + bail!("key column did not come back as Utf8"); + } + (0..batch.num_rows()) + .map(|i| { + Ok(Hit { + key: match keys { + Some(keys) if !keys.is_null(i) => keys.value(i).to_string(), + _ => row_ids.value(i).to_string(), + }, + distance: distances.value(i), + row_id: row_ids.value(i), + }) + }) + .collect() +} + +#[cfg(test)] +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 tempfile::TempDir; + + const DIM: usize = 4; + + fn vector_column(vectors: &[Vec]) -> FixedSizeListArray { + FixedSizeListArray::from_iter_primitive::( + vectors + .iter() + .map(|v| Some(v.iter().copied().map(Some).collect::>())), + DIM as i32, + ) + } + + async fn write_dataset(dir: &TempDir, vectors: &[Vec], 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 = (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::>(), + )), + ], + ) + .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() + } + + #[tokio::test] + async fn flat_returns_exact_top_k() { + let dir = TempDir::new().unwrap(); + let vectors = vec![ + vec![1.0, 0.0, 0.0, 0.0], + vec![0.9, 0.1, 0.0, 0.0], + vec![0.0, 1.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + ]; + let dataset = write_dataset(&dir, &vectors, &["a", "b", "a", "b"]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + + let query = Query { + vector: vec![1.0, 0.0, 0.0, 0.0], + k: 2, + filter: None, + distance: DistanceType::L2, + nprobes: None, + }; + let hits = flat(&dataset, &table, &query).await.unwrap(); + assert_eq!(hits.len(), 2); + assert_eq!(hits[0].key, "k0"); + assert!(hits[0].distance <= hits[1].distance); + + let mut brute: Vec<(String, f32)> = vectors + .iter() + .enumerate() + .map(|(i, v)| (format!("k{i}"), l2(&query.vector, v))) + .collect(); + brute.sort_by(|a, b| a.1.total_cmp(&b.1)); + for (hit, (key, distance)) in hits.iter().zip(brute.iter()) { + assert_eq!(&hit.key, key); + assert!((hit.distance - distance).abs() < 1e-5); + } + } + + #[tokio::test] + async fn flat_applies_the_filter() { + let dir = TempDir::new().unwrap(); + let vectors = vec![ + vec![1.0, 0.0, 0.0, 0.0], + vec![0.9, 0.1, 0.0, 0.0], + vec![0.0, 1.0, 0.0, 0.0], + ]; + let dataset = write_dataset(&dir, &vectors, &["a", "b", "a"]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + + let hits = flat( + &dataset, + &table, + &Query { + vector: vec![1.0, 0.0, 0.0, 0.0], + k: 3, + filter: Some("tag = 'a'".to_string()), + distance: DistanceType::L2, + nprobes: None, + }, + ) + .await + .unwrap(); + assert_eq!(hits.len(), 2); + assert_eq!(hits[0].key, "k0"); + assert_eq!(hits[1].key, "k2"); + } + + #[tokio::test] + async fn get_fetches_rows_by_key() { + let dir = TempDir::new().unwrap(); + let vectors = vec![vec![1.0; DIM], vec![2.0; DIM], vec![3.0; DIM]]; + let dataset = write_dataset(&dir, &vectors, &["a", "b", "c"]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + + let batches = get(&dataset, &table, &["k0".to_string(), "k2".to_string()]) + .await + .unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 2); + } + + #[tokio::test] + async fn query_rejects_a_wrong_dimension() { + let dir = TempDir::new().unwrap(); + let vectors = vec![vec![1.0; DIM]]; + let dataset = write_dataset(&dir, &vectors, &["a"]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + let query = Query { + vector: vec![1.0; DIM + 1], + k: 1, + filter: None, + distance: DistanceType::L2, + nprobes: None, + }; + assert!(flat(&dataset, &table, &query).await.is_err()); + } + + #[tokio::test] + async fn ann_matches_flat_without_an_index() { + let dir = TempDir::new().unwrap(); + let vectors = vec![ + vec![1.0, 0.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + ]; + let dataset = write_dataset(&dir, &vectors, &["a", "b", "c"]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + let query = Query { + vector: vec![1.0, 0.0, 0.0, 0.0], + k: 2, + filter: None, + distance: DistanceType::L2, + nprobes: None, + }; + let indexed = ann(&dataset, &table, &query).await.unwrap(); + let exact = flat(&dataset, &table, &query).await.unwrap(); + assert_eq!( + indexed.iter().map(|h| &h.key).collect::>(), + exact.iter().map(|h| &h.key).collect::>() + ); + } +} diff --git a/seaweed-vector/crates/core/src/table.rs b/seaweed-vector/crates/core/src/table.rs new file mode 100644 index 0000000000..824ccf631e --- /dev/null +++ b/seaweed-vector/crates/core/src/table.rs @@ -0,0 +1,178 @@ +//! The columns a vector query needs, resolved from a table's Arrow schema. + +use anyhow::{Context, Result, bail}; +use arrow_schema::DataType; +use lance::Dataset; + +/// The default name for the vector column, matching what the S3 Vectors API +/// writes. +const DEFAULT_VECTOR_COLUMN: &str = "vector"; + +/// The default name for the key column. +const DEFAULT_KEY_COLUMN: &str = "key"; + +/// Layout of a Lance table used as a vector index. +#[derive(Debug, Clone)] +pub struct VectorTable { + /// The row key column; `None` falls back to `_rowid` in results. + pub key_column: Option, + pub vector_column: String, + pub dimension: usize, + pub metadata_columns: Vec, +} + +impl VectorTable { + /// Resolve the vector layout of `dataset`. `vector_column` and + /// `key_column` override detection: the vector column defaults to the one + /// named `vector`, or the sole fixed-size float list when unambiguous. + pub fn resolve( + dataset: &Dataset, + vector_column: Option<&str>, + key_column: Option<&str>, + ) -> Result { + let schema: arrow_schema::Schema = dataset.schema().into(); + resolve_schema(&schema, vector_column, key_column) + } +} + +pub(crate) fn resolve_schema( + schema: &arrow_schema::Schema, + vector_column: Option<&str>, + key_column: Option<&str>, +) -> Result { + let vector_column = match vector_column { + Some(name) => name.to_string(), + None => detect_vector_column(schema)?, + }; + let field = schema + .field_with_name(&vector_column) + .with_context(|| format!("no such column {vector_column}"))?; + let dimension = vector_dimension(field.data_type()) + .with_context(|| format!("column {vector_column} is not a fixed-size float vector"))?; + + let key_column = match key_column { + Some(name) => { + let field = schema + .field_with_name(name) + .with_context(|| format!("no such column {name}"))?; + if field.data_type() != &DataType::Utf8 { + bail!("key column {name} is not Utf8"); + } + Some(name.to_string()) + } + None => detect_key_column(schema), + }; + + let metadata_columns = schema + .fields() + .iter() + .map(|f| f.name().clone()) + .filter(|name| name != &vector_column && Some(name) != key_column.as_ref()) + .collect(); + + Ok(VectorTable { + key_column, + vector_column, + dimension, + metadata_columns, + }) +} + +fn detect_vector_column(schema: &arrow_schema::Schema) -> Result { + let mut named = None; + let mut candidates = Vec::new(); + for field in schema.fields() { + if vector_dimension(field.data_type()).is_none() { + continue; + } + if field.name() == DEFAULT_VECTOR_COLUMN { + named = Some(field.name().clone()); + } + candidates.push(field.name().clone()); + } + if let Some(name) = named { + return Ok(name); + } + match candidates.as_slice() { + [only] => Ok(only.clone()), + [] => bail!("no fixed-size float vector column found"), + many => bail!("several vector columns, pick one of: {}", many.join(", ")), + } +} + +fn detect_key_column(schema: &arrow_schema::Schema) -> Option { + schema + .field_with_name(DEFAULT_KEY_COLUMN) + .ok() + .filter(|f| f.data_type() == &DataType::Utf8) + .map(|f| f.name().clone()) +} + +fn vector_dimension(data_type: &DataType) -> Option { + let DataType::FixedSizeList(child, dimension) = data_type else { + return None; + }; + match child.data_type() { + DataType::Float16 | DataType::Float32 | DataType::Float64 => Some(*dimension as usize), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::{Field, Schema}; + use std::sync::Arc; + + fn vector_field(name: &str, dim: i32) -> Field { + Field::new( + name, + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ) + } + + #[test] + fn resolve_prefers_the_column_named_vector() { + let schema = Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + vector_field("embedding", 4), + vector_field("vector", 8), + Field::new("tag", DataType::Utf8, true), + ]); + let table = resolve_schema(&schema, None, None).unwrap(); + assert_eq!(table.vector_column, "vector"); + assert_eq!(table.dimension, 8); + assert_eq!(table.key_column.as_deref(), Some("key")); + assert_eq!(table.metadata_columns, vec!["embedding", "tag"]); + } + + #[test] + fn resolve_uses_the_sole_vector_column() { + let schema = Schema::new(vec![ + vector_field("embedding", 4), + Field::new("tag", DataType::Utf8, true), + ]); + let table = resolve_schema(&schema, None, None).unwrap(); + assert_eq!(table.vector_column, "embedding"); + assert_eq!(table.key_column, None); + } + + #[test] + fn resolve_rejects_ambiguous_vector_columns() { + let schema = Schema::new(vec![vector_field("a", 4), vector_field("b", 4)]); + assert!(resolve_schema(&schema, None, None).is_err()); + } + + #[test] + fn resolve_rejects_a_non_vector_column() { + let schema = Schema::new(vec![ + Field::new("vector", DataType::Utf8, true), + vector_field("embedding", 4), + ]); + assert!(resolve_schema(&schema, Some("vector"), None).is_err()); + let table = resolve_schema(&schema, None, Some("vector")).unwrap(); + assert_eq!(table.vector_column, "embedding"); + assert_eq!(table.key_column.as_deref(), Some("vector")); + } +}