vector: seaweed-vector workspace + flat/indexed top-K engine (P0 phase 1) (#12096)

* 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

@@ -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"] }

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

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;

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<f32>,
+    pub k: usize,
+    /// A Lance SQL filter over the metadata columns.
+    pub filter: Option<String>,
+    pub distance: DistanceType,
+    pub nprobes: Option<usize>,
+}
+
+/// 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<Vec<Hit>> {
+    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<Vec<Hit>> {
+    run(scanner(dataset, table, query, true)?, table).await
+}
+
+/// Fetch rows by key.
+pub async fn get(
+    dataset: &Dataset,
+    table: &VectorTable,
+    keys: &[String],
+) -> Result<Vec<RecordBatch>> {
+    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::<Vec<_>>()
+        .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<Scanner> {
+    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<Vec<Hit>> {
+    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<Vec<Hit>> {
+    let distances = batch
+        .column_by_name(DISTANCE_COLUMN)
+        .and_then(|c| c.as_any().downcast_ref::<Float32Array>())
+        .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::<UInt64Array>())
+        .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::<StringArray>())
+    });
+    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<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()
+    }
+
+    #[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::<Vec<_>>(),
+            exact.iter().map(|h| &h.key).collect::<Vec<_>>()
+        );
+    }
+}

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<String>,
+    pub vector_column: String,
+    pub dimension: usize,
+    pub metadata_columns: Vec<String>,
+}
+
+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<Self> {
+        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<VectorTable> {
+    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<String> {
+    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<String> {
+    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<usize> {
+    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"));
+    }
+}