diff --git a/.github/workflows/rust-vector-tests.yml b/.github/workflows/rust-vector-tests.yml new file mode 100644 index 0000000000..ca497ee71e --- /dev/null +++ b/.github/workflows/rust-vector-tests.yml @@ -0,0 +1,76 @@ +name: "Rust Vector Tests" + +on: + pull_request: + branches: [ master ] + paths: + - 'seaweed-vector/**' + - 'seaweed-worker/**' + - 'seaweed-common/**' + - 'weed/pb/plugin.proto' + - '.github/workflows/rust-vector-tests.yml' + push: + branches: [ master, main ] + paths: + - 'seaweed-vector/**' + - 'seaweed-worker/**' + - 'seaweed-common/**' + - 'weed/pb/plugin.proto' + - '.github/workflows/rust-vector-tests.yml' + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || (github.event_name == 'push' && github.ref) || github.run_id }} + cancel-in-progress: ${{ github.event_name == 'pull_request' || (github.event_name == 'push' && !startsWith(github.ref, 'refs/tags/')) }} + +permissions: + contents: read + +jobs: + rust-vector-build: + name: Rust Vector Build and Unit Tests + runs-on: ubuntu-24.04 + timeout-minutes: 45 + + steps: + - name: Checkout code + uses: actions/checkout@v7 + with: + persist-credentials: false + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + + - name: Fingerprint build toolchain + id: toolchain + run: echo "fingerprint=$(getconf GNU_LIBC_VERSION | tr ' ' '-')-rustc-$(rustc -V | awk '{print $2}')" >> "$GITHUB_OUTPUT" + + - name: Cache cargo registry and target + if: runner.environment == 'github-hosted' + uses: actions/cache@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + seaweed-vector/target/release + key: rust-vector-${{ steps.toolchain.outputs.fingerprint }}-${{ hashFiles('seaweed-vector/Cargo.lock') }} + restore-keys: | + rust-vector-${{ steps.toolchain.outputs.fingerprint }}- + + # lance's build scripts compile their own protos and look for a protoc. + # Point them at the one protoc-bin-vendored ships, so no job depends on a + # system package and every build sees the same version. + - name: Use the vendored protoc + run: | + cd seaweed-vector + cargo fetch + version=$(awk '/^name = "protoc-bin-vendored-linux-x86_64"$/{found=1; next} found && /^version = /{gsub(/"/,"",$3); print $3; exit}' Cargo.lock) + test -n "$version" || { echo "protoc-bin-vendored-linux-x86_64 is not in Cargo.lock" >&2; exit 1; } + protoc=$(find ~/.cargo/registry/src -path "*protoc-bin-vendored-linux-x86_64-$version/bin/protoc" | head -1) + test -x "$protoc" || { echo "no vendored protoc $version in the registry" >&2; exit 1; } + echo "PROTOC=$protoc" >> "$GITHUB_ENV" + + - name: Build the vector crates + run: cd seaweed-vector && cargo build --release + + - name: Run unit tests + run: cd seaweed-vector && cargo test --release --workspace diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index c4780cdd26..879a1c4fbb 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6457,7 +6457,9 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow-array", + "arrow-cast", "arrow-schema", + "clap", "futures", "lance", "lance-index", @@ -6469,6 +6471,7 @@ dependencies = [ "tempfile", "tokio", "tracing", + "tracing-subscriber", "weed-lance-worker", ] diff --git a/seaweed-vector/crates/core/Cargo.toml b/seaweed-vector/crates/core/Cargo.toml index 4241f4a136..7888bbe51b 100644 --- a/seaweed-vector/crates/core/Cargo.toml +++ b/seaweed-vector/crates/core/Cargo.toml @@ -8,6 +8,10 @@ description = "Vector query engine over SeaweedFS Lance table buckets" [lib] name = "weed_vector_core" +[[bin]] +name = "weed-vector" +path = "src/main.rs" + [dependencies] # The catalog client and dataset-opening helper the maintenance worker uses. # A path dependency because the two trees are separate cargo workspaces with @@ -18,14 +22,17 @@ lance-index = "10" lance-linalg = "10" lance-table = "10" arrow-array = "58" +arrow-cast = "58" arrow-schema = "58" roaring = "0.11" anyhow.workspace = true +clap = { version = "4", features = ["derive", "env"] } futures = "0.3" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio.workspace = true tracing.workspace = true +tracing-subscriber.workspace = true [dev-dependencies] tempfile = "3" diff --git a/seaweed-vector/crates/core/src/coverage.rs b/seaweed-vector/crates/core/src/coverage.rs index 02e537922f..4dd68620a6 100644 --- a/seaweed-vector/crates/core/src/coverage.rs +++ b/seaweed-vector/crates/core/src/coverage.rs @@ -8,7 +8,7 @@ use lance::index::DatasetIndexExt; use roaring::RoaringBitmap; /// How much of the dataset one named index covers at this version. -#[derive(Debug)] +#[derive(Debug, serde::Serialize)] pub struct IndexCoverage { pub index_name: String, /// Physical index segments: the IndexMetadata deltas sharing the name. diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs index c4842ddbae..e4cb1453d2 100644 --- a/seaweed-vector/crates/core/src/lib.rs +++ b/seaweed-vector/crates/core/src/lib.rs @@ -10,6 +10,8 @@ pub mod query; pub mod table; #[cfg(test)] mod testing; +#[cfg(test)] +mod verify; pub use coverage::IndexCoverage; pub use query::{Hit, Query}; diff --git a/seaweed-vector/crates/core/src/main.rs b/seaweed-vector/crates/core/src/main.rs new file mode 100644 index 0000000000..879dfb362a --- /dev/null +++ b/seaweed-vector/crates/core/src/main.rs @@ -0,0 +1,209 @@ +use anyhow::{Context, Result, bail}; +use clap::{Parser, Subcommand}; + +use weed_lance_worker::catalog::{NamespaceClient, parse_id}; +use weed_lance_worker::dataset::{self, FallbackOptions}; +use weed_vector_core::coverage; +use weed_vector_core::query::{self, Query}; +use weed_vector_core::table::VectorTable; + +/// Mirrors `weed-worker`'s flags where the two touch the same thing: the +/// namespace is WEED_LANCE_NAMESPACE and the fallback storage credentials are +/// the same env vars. +#[derive(Parser, Debug)] +#[command(name = "weed-vector", about = "SeaweedFS vector query spike")] +struct Args { + /// Lance namespace the tables live in. + #[arg( + long, + default_value = "http://localhost:9101", + env = "WEED_LANCE_NAMESPACE" + )] + namespace: String, + + /// Storage credentials to use where the namespace vends none. + #[arg(long, env = "WEED_S3_ACCESS_KEY")] + access_key: Option, + + #[arg(long, env = "WEED_S3_SECRET_KEY")] + secret_key: Option, + + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand, Debug)] +enum Command { + /// A table's location, version, vector layout and index coverage. + Describe { + /// Delimiter-joined table id, e.g. bucket$table. + id: String, + }, + /// Which fragments each vector index covers. + Coverage { id: String }, + /// A top-K query over the whole dataset. + Query { + id: String, + /// Comma-separated floats, e.g. --vector 0.1,0.2,0.3,0.4 + #[arg(long)] + vector: String, + #[arg(long, default_value = "10")] + k: usize, + /// Lance SQL filter over the metadata columns. + #[arg(long)] + filter: Option, + /// l2, cosine or dot. + #[arg(long, default_value = "l2")] + metric: String, + /// Exact scan, no index. + #[arg(long)] + flat: bool, + /// Filter the top-K rows after ranking (Lance's default); off keeps + /// the filter before the ANN search. + #[arg(long)] + postfilter: bool, + /// Pin the query to this committed version. + #[arg(long)] + version: Option, + #[arg(long)] + vector_column: Option, + #[arg(long)] + key_column: Option, + #[arg(long)] + nprobes: Option, + }, + /// Fetch rows by key. + Get { id: String, keys: Vec }, +} + +/// `id` is a delimiter-joined table id; when it looks like a URI (`s3://...`, +/// `/path/...`, `file://...`) the dataset is opened directly, skipping the +/// namespace — useful for a spike run against a local dataset. +async fn open_dataset( + client: &NamespaceClient, + id: &str, + fallback: &FallbackOptions, +) -> Result { + if id.contains("://") || id.starts_with('/') { + let mut builder = lance::dataset::builder::DatasetBuilder::from_uri(id); + if !fallback.is_empty() { + builder = builder.with_storage_options(fallback.clone()); + } + return builder.load().await.with_context(|| format!("open {id}")); + } + let table = dataset::open(client, &parse_id(id), fallback).await?; + Ok(table.dataset) +} + +#[tokio::main] +async fn main() -> Result<()> { + weed_lance_worker::tls::install_default_crypto_provider(); + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + let args = Args::parse(); + + let mut fallback = FallbackOptions::new(); + if let (Some(access), Some(secret)) = (args.access_key, args.secret_key) { + fallback.insert("aws_access_key_id".to_string(), access); + fallback.insert("aws_secret_access_key".to_string(), secret); + } + let client = NamespaceClient::new(args.namespace).with_fallback(&fallback); + + match args.command { + Command::Describe { id } => { + let dataset = open_dataset(&client, &id, &fallback).await?; + let schema: arrow_schema::Schema = dataset.schema().into(); + let vector = VectorTable::resolve(&dataset, None, None).ok(); + let coverages = match &vector { + Some(v) => coverage::coverage(&dataset, &v.vector_column).await?, + None => Vec::new(), + }; + println!( + "{}", + serde_json::json!({ + "id": id, + "uri": dataset.uri(), + "version": dataset.version().version, + "fragments": dataset.get_fragments().len(), + "rows": dataset.count_rows(None).await?, + "schema": schema.fields().iter().map(|f| serde_json::json!({ + "name": f.name(), + "type": f.data_type().to_string(), + "nullable": f.is_nullable(), + })).collect::>(), + "vector": vector.map(|v| serde_json::json!({ + "vector_column": v.vector_column, + "dimension": v.dimension, + "key_column": v.key_column, + "metadata_columns": v.metadata_columns, + })), + "coverage": coverages, + }) + ); + } + Command::Coverage { id } => { + let dataset = open_dataset(&client, &id, &fallback).await?; + let vector = VectorTable::resolve(&dataset, None, None)?; + let coverages = coverage::coverage(&dataset, &vector.vector_column).await?; + println!("{}", serde_json::to_string_pretty(&coverages)?); + } + Command::Get { id, keys } => { + let dataset = open_dataset(&client, &id, &fallback).await?; + let vector = VectorTable::resolve(&dataset, None, None)?; + for batch in query::get(&dataset, &vector, &keys).await? { + println!("{}", arrow_cast::pretty::pretty_format_batches(&[batch])?); + } + } + Command::Query { + id, + vector, + k, + filter, + metric, + flat, + postfilter, + version, + vector_column, + key_column, + nprobes, + } => { + let opened = open_dataset(&client, &id, &fallback).await?; + let dataset = match version { + Some(v) => opened.checkout_version(v).await?, + None => opened, + }; + let vector_table = + VectorTable::resolve(&dataset, vector_column.as_deref(), key_column.as_deref())?; + let vector: Vec = vector + .split(',') + .map(|part| part.trim().parse::()) + .collect::>() + .context("--vector takes comma-separated floats")?; + let distance = match metric.as_str() { + "l2" | "euclidean" => lance_linalg::distance::DistanceType::L2, + "cosine" => lance_linalg::distance::DistanceType::Cosine, + "dot" | "dotproduct" => lance_linalg::distance::DistanceType::Dot, + other => bail!("unknown metric {other}: l2, cosine or dot"), + }; + let query = Query { + vector, + k, + filter, + postfilter, + distance, + nprobes, + }; + let hits = if flat { + query::flat(&dataset, &vector_table, &query).await? + } else { + query::ann(&dataset, &vector_table, &query).await? + }; + println!("{}", serde_json::to_string_pretty(&hits)?); + } + } + Ok(()) +} diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs index fb917adfec..420afb69fb 100644 --- a/seaweed-vector/crates/core/src/query.rs +++ b/seaweed-vector/crates/core/src/query.rs @@ -30,7 +30,7 @@ pub struct Query { } /// One row of a top-K answer. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize)] pub struct Hit { /// The key column value, or the row id when the table has no key column. pub key: String, diff --git a/seaweed-vector/crates/core/src/verify.rs b/seaweed-vector/crates/core/src/verify.rs new file mode 100644 index 0000000000..2ba0b91b8d --- /dev/null +++ b/seaweed-vector/crates/core/src/verify.rs @@ -0,0 +1,201 @@ +//! The spike's exit gate: the engine's answers must equal what a bare Lance +//! `scan().nearest()` produces on the same pinned version, and no eligible +//! row — indexed or not — may be missing. + +use arrow_array::{Float32Array, UInt64Array}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::index::DatasetIndexExt; +use lance::index::vector::VectorIndexParams; +use lance_index::IndexType; +use lance_index::optimize::OptimizeOptions; +use lance_linalg::distance::DistanceType; +use tempfile::TempDir; + +use crate::coverage::coverage; +use crate::query::{self, Query}; +use crate::table::VectorTable; +use crate::testing::{self, DIM}; + +async fn indexed_dataset(dir: &TempDir) -> Dataset { + let mut dataset = testing::write_dataset(dir, &testing::vectors(64, DIM), &["a"; 64]).await; + let params = VectorIndexParams::ivf_flat(4, DistanceType::L2); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + dataset +} + +fn query_for(vector: Vec, k: usize) -> Query { + Query { + vector, + k, + filter: None, + postfilter: false, + distance: DistanceType::L2, + nprobes: Some(4), + } +} + +/// The reference path: a bare scanner over the same dataset handle. +async fn reference(dataset: &Dataset, vector: &[f32], k: usize) -> Vec<(u64, f32)> { + let mut scanner = dataset.scan(); + scanner + .nearest("vector", &Float32Array::from(vector.to_vec()), k) + .unwrap() + .distance_metric(DistanceType::L2) + .nprobes(4) + .with_row_id(); + let mut stream = scanner.try_into_stream().await.unwrap(); + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await.unwrap() { + let ids = batch + .column_by_name("_rowid") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let distances = batch + .column_by_name("_distance") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + rows.push((ids.value(i), distances.value(i))); + } + } + rows +} + +async fn assert_matches_reference(dataset: &Dataset, vector: &[f32], k: usize) { + let table = VectorTable::resolve(dataset, None, None).unwrap(); + let hits = query::ann(dataset, &table, &query_for(vector.to_vec(), k)) + .await + .unwrap(); + let expected = reference(dataset, vector, k).await; + assert_eq!(hits.len(), expected.len()); + for (hit, (row_id, distance)) in hits.iter().zip(expected.iter()) { + assert_eq!(hit.row_id, *row_id); + assert!((hit.distance - distance).abs() < 1e-6, "{}", hit.key); + } +} + +#[tokio::test] +async fn ann_matches_lance_reference() { + let dir = TempDir::new().unwrap(); + let dataset = indexed_dataset(&dir).await; + assert_matches_reference(&dataset, &testing::vectors(1, DIM)[0], 10).await; +} + +#[tokio::test] +async fn ann_matches_reference_over_unindexed_fragments() { + let dir = TempDir::new().unwrap(); + let mut dataset = indexed_dataset(&dir).await; + testing::append_dataset(&mut dataset, &testing::vectors(8, DIM), &["a"; 8]).await; + + let cov = &coverage(&dataset, "vector").await.unwrap()[0]; + assert_eq!(cov.unindexed_fragments.len(), 1); + assert_eq!(cov.unindexed_rows, 8); + + for vector in testing::vectors(4, DIM) { + assert_matches_reference(&dataset, &vector, 10).await; + } +} + +#[tokio::test] +async fn ann_matches_reference_across_segments() { + let dir = TempDir::new().unwrap(); + let mut dataset = indexed_dataset(&dir).await; + testing::append_dataset(&mut dataset, &testing::vectors(32, DIM), &["a"; 32]).await; + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let cov = &coverage(&dataset, "vector").await.unwrap()[0]; + assert_eq!(cov.segments, 2); + assert!(cov.unindexed_fragments.is_empty()); + + for vector in testing::vectors(4, DIM) { + assert_matches_reference(&dataset, &vector, 10).await; + } +} + +#[tokio::test] +async fn deletes_are_not_returned() { + let dir = TempDir::new().unwrap(); + let mut dataset = indexed_dataset(&dir).await; + dataset.delete("key = 'k0'").await.unwrap(); + assert_matches_reference(&dataset, &testing::vectors(1, DIM)[0], 10).await; +} + +/// Writes a small indexed dataset with unindexed rows to +/// $WEED_VECTOR_SPIKE_DIR for exercising the `weed-vector` CLI: +/// +/// WEED_VECTOR_SPIKE_DIR=/tmp/vds cargo test make_spike_dataset -- --ignored +/// weed-vector describe /tmp/vds +/// weed-vector coverage /tmp/vds +/// weed-vector query /tmp/vds --vector 0,0,0,0 --k 5 +#[tokio::test] +#[ignore] +async fn make_spike_dataset() { + let dir = std::env::var("WEED_VECTOR_SPIKE_DIR").expect("set WEED_VECTOR_SPIKE_DIR"); + let mut dataset = testing::write_dataset_at(&dir, &testing::vectors(64, DIM), &["a"; 64]).await; + let params = VectorIndexParams::ivf_flat(4, DistanceType::L2); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + testing::append_dataset(&mut dataset, &testing::vectors(8, DIM), &["b"; 8]).await; +} + +#[tokio::test] +async fn query_at_pinned_version_survives_maintenance() { + let dir = TempDir::new().unwrap(); + let mut dataset = indexed_dataset(&dir).await; + let pinned = dataset + .checkout_version(dataset.version().version) + .await + .unwrap(); + + testing::append_dataset(&mut dataset, &testing::vectors(8, DIM), &["b"; 8]).await; + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + + assert_eq!(pinned.count_rows(None).await.unwrap(), 64); + let vector = testing::vectors(1, DIM)[0].clone(); + let hits = query::ann( + &pinned, + &VectorTable::resolve(&pinned, None, None).unwrap(), + &query_for(vector.clone(), 10), + ) + .await + .unwrap(); + let expected = reference(&pinned, &vector, 10).await; + assert_eq!(hits.len(), expected.len()); + assert!(hits.iter().all(|h| h.row_id < 64)); + for (hit, (row_id, distance)) in hits.iter().zip(expected.iter()) { + assert_eq!(hit.row_id, *row_id); + assert!((hit.distance - distance).abs() < 1e-6); + } + + assert_matches_reference(&dataset, &vector, 10).await; +}