diff --git a/.gitignore b/.gitignore index 9bce980521..dc858fdb1e 100644 --- a/.gitignore +++ b/.gitignore @@ -148,5 +148,7 @@ __pycache__/ /test/s3vectors/weed /test/s3vectors/weed-vector /test/s3vectors/weed-vector-worker +/test/s3vectors/weed-vector-cli +/test/s3vectors/weed-worker /test/s3vectors/weed-volume .pytest_cache/ diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index 3891b8f226..2d13f80da6 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6593,6 +6593,7 @@ dependencies = [ "lance", "lance-datafusion", "lance-index", + "lance-linalg", "lance-table", "prometheus", "reqwest 0.12.28", diff --git a/seaweed-vector/crates/core/src/index.rs b/seaweed-vector/crates/core/src/index.rs new file mode 100644 index 0000000000..ca19d0c8b9 --- /dev/null +++ b/seaweed-vector/crates/core/src/index.rs @@ -0,0 +1,238 @@ +//! Sizing the vector index an S3 Vectors table searches through, and the +//! probes a query spends on it. Building lives in the worker crate. + +pub use weed_lance_worker::vector_index::{ + INDEX_NAME, MAX_PARTITIONS, build, distance_metric, index_names, optimize, partitions_for, +}; + +pub const MIN_PROBES: usize = 8; + +/// Partitions one query probes: a tenth of them, at least [`MIN_PROBES`], never +/// more than exist. +pub fn probes_for(partitions: usize) -> usize { + MIN_PROBES + .max(partitions.div_ceil(10)) + .min(partitions.max(1)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::coverage::coverage; + use crate::query::{self, Query}; + use crate::table::VectorTable; + use crate::testing::{self, DIM}; + use lance::Dataset; + use lance::dataset::optimize::{CompactionOptions, compact_files}; + use lance_linalg::distance::DistanceType; + use tempfile::TempDir; + + #[test] + fn probes_are_a_tenth_with_a_floor() { + assert_eq!(probes_for(1), 1); + assert_eq!(probes_for(8), 8); + assert_eq!(probes_for(64), 8); + assert_eq!(probes_for(224), 23); + assert_eq!(probes_for(4096), 410); + } + + fn exact(vector: &[f32], k: usize) -> Query { + Query { + vector: vector.to_vec(), + k, + filter: None, + postfilter: false, + distance: DistanceType::L2, + nprobes: None, + index_segments: None, + fragments: None, + } + } + + /// The test vectors repeat, so ties are common: compare distances, not rows. + async fn assert_same_top_k(dataset: &Dataset, vector: &[f32], k: usize) { + let table = VectorTable::resolve(dataset, None, None).unwrap(); + let q = exact(vector, k); + let flat = query::flat(dataset, &table, &q).await.unwrap(); + let ann = query::ann( + dataset, + &table, + &Query { + nprobes: Some(MAX_PARTITIONS), + ..q + }, + ) + .await + .unwrap(); + assert_eq!(flat.len(), ann.len()); + for (f, a) in flat.iter().zip(ann.iter()) { + assert!( + (f.distance - a.distance).abs() < 1e-6, + "{flat:?} vs {ann:?}" + ); + } + } + + #[tokio::test] + async fn build_covers_every_fragment() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(200, DIM), &["a"; 200]).await; + testing::append_dataset(&mut dataset, &testing::vectors(50, DIM), &["b"; 50]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + build( + &mut dataset, + &table.vector_column, + INDEX_NAME, + partitions_for(250), + false, + ) + .await + .unwrap(); + let cov = coverage(&dataset, "vector").await.unwrap(); + assert_eq!(cov.len(), 1); + assert_eq!(cov[0].index_name, INDEX_NAME); + assert_eq!(cov[0].segments, 1); + assert!(cov[0].coverage_known); + assert!(cov[0].unindexed_fragments.is_empty()); + for vector in testing::vectors(4, DIM) { + assert_same_top_k(&dataset, &vector, 10).await; + } + } + + #[tokio::test] + async fn build_uses_the_tables_metric() { + for (metric, distance) in [ + ("COSINE", DistanceType::Cosine), + ("DOT_PRODUCT", DistanceType::Dot), + ("L2", DistanceType::L2), + ] { + let dir = TempDir::new().unwrap(); + let rows = testing::vectors(300, DIM); + let mut dataset = + testing::write_dataset_with_metric(&dir, &rows, &["a"; 300], metric).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + assert_eq!(distance_metric(&dataset, "vector"), distance); + build(&mut dataset, &table.vector_column, INDEX_NAME, 4, false) + .await + .unwrap(); + for vector in testing::vectors(3, DIM) { + let q = Query { + distance, + ..exact(&vector, 10) + }; + let flat = query::flat(&dataset, &table, &q).await.unwrap(); + let ann = query::ann( + &dataset, + &table, + &Query { + nprobes: Some(4), + ..q + }, + ) + .await + .unwrap(); + assert_eq!(flat.len(), ann.len(), "{metric}"); + for (f, a) in flat.iter().zip(ann.iter()) { + assert!( + (f.distance - a.distance).abs() < 1e-5, + "{metric}: {flat:?} vs {ann:?}" + ); + } + } + } + } + + #[tokio::test] + async fn optimize_appends_then_retrain_folds_segments() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(100, DIM), &["a"; 100]).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + build(&mut dataset, &table.vector_column, INDEX_NAME, 4, false) + .await + .unwrap(); + let mut appended = testing::vectors(40, DIM); + appended.push(vec![9.0; DIM]); + testing::append_dataset(&mut dataset, &appended, &["b"; 41]).await; + optimize(&mut dataset, &table.vector_column, false) + .await + .unwrap(); + let cov = &coverage(&dataset, "vector").await.unwrap()[0]; + assert_eq!(cov.segments, 2); + assert!(cov.unindexed_fragments.is_empty()); + let hits = query::ann( + &dataset, + &table, + &Query { + nprobes: Some(4), + ..exact(&[9.0; DIM], 1) + }, + ) + .await + .unwrap(); + assert_eq!(hits[0].key, "k140", "{hits:?}"); + + testing::append_dataset(&mut dataset, &testing::vectors(40, DIM), &["c"; 40]).await; + optimize(&mut dataset, &table.vector_column, true) + .await + .unwrap(); + let cov = &coverage(&dataset, "vector").await.unwrap()[0]; + assert_eq!(cov.segments, 1); + assert!(cov.coverage_known); + assert!(cov.unindexed_fragments.is_empty()); + for vector in testing::vectors(4, DIM) { + assert_same_top_k(&dataset, &vector, 10).await; + } + } + + #[tokio::test] + async fn compaction_keeps_the_index_covering() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(60, DIM), &["a"; 60]).await; + for tag in ["b", "c", "d"] { + testing::append_dataset(&mut dataset, &testing::vectors(60, DIM), &[tag; 60]).await; + } + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + build(&mut dataset, &table.vector_column, INDEX_NAME, 4, false) + .await + .unwrap(); + dataset.delete("key = 'k0'").await.unwrap(); + assert_eq!(dataset.get_fragments().len(), 4); + + let options = CompactionOptions { + target_rows_per_fragment: 1000, + materialize_deletions: true, + materialize_deletions_threshold: 0.0, + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.get_fragments().len(), 1); + + let cov = &coverage(&dataset, "vector").await.unwrap()[0]; + assert!(cov.coverage_known); + assert!(cov.unindexed_fragments.is_empty(), "{cov:?}"); + let hits = query::ann( + &dataset, + &table, + &Query { + nprobes: Some(MAX_PARTITIONS), + ..exact(&testing::vectors(1, DIM)[0], 10) + }, + ) + .await + .unwrap(); + let mut tied: Vec<&str> = hits + .iter() + .filter(|h| h.distance < 1e-6) + .map(|h| h.key.as_str()) + .collect(); + tied.sort(); + assert_eq!(tied, ["k120", "k180", "k60"], "{hits:?}"); + for vector in testing::vectors(4, DIM) { + assert_same_top_k(&dataset, &vector, 10).await; + } + } +} diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs index 5d8013ed24..be63ae04b2 100644 --- a/seaweed-vector/crates/core/src/lib.rs +++ b/seaweed-vector/crates/core/src/lib.rs @@ -7,6 +7,7 @@ pub mod coverage; mod filter; +pub mod index; pub mod maintenance; pub mod query; pub mod segments; diff --git a/seaweed-vector/crates/core/src/main.rs b/seaweed-vector/crates/core/src/main.rs index cc293f3067..6299d64740 100644 --- a/seaweed-vector/crates/core/src/main.rs +++ b/seaweed-vector/crates/core/src/main.rs @@ -4,6 +4,7 @@ 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::index; use weed_vector_core::query::{self, Query}; use weed_vector_core::table::VectorTable; @@ -91,6 +92,38 @@ enum Command { #[arg(long)] key_column: Option, }, + /// Build or extend the vector index, then print its coverage. + Index { + #[command(subcommand)] + action: IndexAction, + }, +} + +#[derive(Subcommand, Debug)] +enum IndexAction { + /// An IVF_FLAT index in the table's metric; partitions default to the + /// square root of the row count. + Build { + id: String, + #[arg(long)] + partitions: Option, + #[arg(long, default_value = index::INDEX_NAME)] + name: String, + /// Rebuild an index of that name instead of failing. + #[arg(long)] + replace: bool, + #[arg(long)] + vector_column: Option, + }, + /// Cover the rows written since the index was built. + Optimize { + id: String, + /// Train the index again on every row instead of adding a segment. + #[arg(long)] + retrain: bool, + #[arg(long)] + vector_column: Option, + }, } /// `id` is a delimiter-joined table id; when it looks like a URI (`s3://...`, @@ -116,6 +149,7 @@ async fn open_dataset( async fn main() -> Result<()> { weed_lance_worker::tls::install_default_crypto_provider(); tracing_subscriber::fmt() + .with_writer(std::io::stderr) .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), @@ -243,6 +277,49 @@ async fn main() -> Result<()> { }; println!("{}", serde_json::to_string_pretty(&hits)?); } + Command::Index { action } => { + let (id, vector_column) = match &action { + IndexAction::Build { + id, vector_column, .. + } + | IndexAction::Optimize { + id, vector_column, .. + } => (id.clone(), vector_column.clone()), + }; + let mut dataset = open_dataset(&client, &id, &fallback).await?; + let table = VectorTable::resolve(&dataset, vector_column.as_deref(), None)?; + let started = std::time::Instant::now(); + match action { + IndexAction::Build { + partitions, + name, + replace, + .. + } => { + let rows = dataset.count_rows(None).await?; + let partitions = partitions.unwrap_or_else(|| index::partitions_for(rows)); + index::build( + &mut dataset, + &table.vector_column, + &name, + partitions, + replace, + ) + .await?; + } + IndexAction::Optimize { retrain, .. } => { + index::optimize(&mut dataset, &table.vector_column, retrain).await?; + } + } + println!( + "{}", + serde_json::json!({ + "version": dataset.version().version, + "seconds": started.elapsed().as_secs_f64(), + "coverage": coverage::coverage(&dataset, &table.vector_column).await?, + }) + ); + } } Ok(()) } @@ -272,6 +349,38 @@ mod tests { } } + #[test] + fn index_build_takes_partitions_and_a_name() { + let args = Args::try_parse_from([ + "weed-vector", + "index", + "build", + "b$vector$i", + "--partitions", + "32", + "--replace", + ]) + .unwrap(); + match args.command { + Command::Index { + action: + IndexAction::Build { + id, + partitions, + name, + replace, + .. + }, + } => { + assert_eq!(id, "b$vector$i"); + assert_eq!(partitions, Some(32)); + assert_eq!(name, index::INDEX_NAME); + assert!(replace); + } + _ => panic!("expected index build"), + } + } + #[test] fn column_flags_on_get_describe_coverage() { for argv in [ diff --git a/seaweed-vector/crates/core/src/testing.rs b/seaweed-vector/crates/core/src/testing.rs index 9ed419fca7..7f2e9e6043 100644 --- a/seaweed-vector/crates/core/src/testing.rs +++ b/seaweed-vector/crates/core/src/testing.rs @@ -66,6 +66,40 @@ pub async fn write_dataset_at(uri: &str, vectors: &[Vec], tags: &[&str]) -> Dataset::write(reader, uri, None).await.unwrap() } +/// Like [`write_dataset`], with the vector field carrying the `distance_metric` +/// metadata `CreateIndex` writes (`L2`, `COSINE`, `DOT_PRODUCT`). +pub async fn write_dataset_with_metric( + dir: &TempDir, + vectors: &[Vec], + tags: &[&str], + metric: &str, +) -> Dataset { + let dim = vectors[0].len(); + let plain = schema(dim); + let fields: Vec> = + plain + .fields() + .iter() + .map(|f| { + if f.name() == "vector" { + Arc::new(f.as_ref().clone().with_metadata( + [("distance_metric".to_string(), metric.to_string())].into(), + )) + } else { + f.clone() + } + }) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let batch = batch(0, vectors, tags, dim) + .with_schema(schema.clone()) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + Dataset::write(reader, dir.path().to_str().unwrap(), None) + .await + .unwrap() +} + /// Deterministic non-trivial vectors for index coverage tests. pub fn vectors(rows: usize, dim: usize) -> Vec> { (0..rows) diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index 8ccf041c2b..66d2c44d9e 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -12,13 +12,13 @@ use lance::io::ObjectStoreParams; use lance::io::StorageOptionsAccessor; use lance::io::WrappingObjectStore; use lance::table::format::Fragment; -use lance_linalg::distance::DistanceType; use prost_types::{ListValue, Struct, Value, value::Kind}; use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use tonic::{Request, Response, Status}; use weed_lance_worker::catalog::NamespaceClient; use weed_lance_worker::dataset::{self, FallbackOptions}; +use weed_vector_core::index; use weed_vector_core::query::{self, Query}; use weed_vector_core::table::VectorTable; @@ -658,16 +658,7 @@ impl VectorService for VectorServer { if req.top_k == 0 { return Err(Status::invalid_argument("top_k must be positive")); } - // The index's configured metric lives in the vector field's metadata. - let distance = match dataset - .schema() - .field(VECTOR_COLUMN) - .and_then(|f| f.metadata.get("distance_metric").map(String::as_str)) - { - Some("COSINE") => DistanceType::Cosine, - Some("DOT_PRODUCT") => DistanceType::Dot, - _ => DistanceType::L2, - }; + let distance = index::distance_metric(&dataset, VECTOR_COLUMN); let query = Query { vector: req.query_vector, k: req.top_k as usize, diff --git a/seaweed-worker/crates/lance/Cargo.toml b/seaweed-worker/crates/lance/Cargo.toml index 1637ab804f..0ce5f4c9cd 100644 --- a/seaweed-worker/crates/lance/Cargo.toml +++ b/seaweed-worker/crates/lance/Cargo.toml @@ -27,6 +27,7 @@ prometheus.workspace = true # build time this worker never spends. lance = { version = "10", default-features = false, features = ["aws"] } lance-index = "10" +lance-linalg = "10" # For the manifest's own fragment type, which the sort marker summarises. lance-table = "10" # For the execution options the sort needs; the scanner's own stream helper diff --git a/seaweed-worker/crates/lance/src/lib.rs b/seaweed-worker/crates/lance/src/lib.rs index 69455e2aa0..0e69a16447 100644 --- a/seaweed-worker/crates/lance/src/lib.rs +++ b/seaweed-worker/crates/lance/src/lib.rs @@ -11,5 +11,6 @@ pub mod jobs; pub mod metrics; pub mod preview; pub mod tls; +pub mod vector_index; pub use jobs::handlers; diff --git a/seaweed-worker/crates/lance/src/vector_index.rs b/seaweed-worker/crates/lance/src/vector_index.rs new file mode 100644 index 0000000000..f3356273ae --- /dev/null +++ b/seaweed-worker/crates/lance/src/vector_index.rs @@ -0,0 +1,105 @@ +//! The vector index an S3 Vectors table searches through: how it is named, +//! sized and built. The query engine and the maintenance worker share this. + +use anyhow::{Context, Result}; +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; + +/// The name every S3 Vectors index build uses, so a table has one logical index. +pub const INDEX_NAME: &str = "vector_idx"; +pub const MAX_PARTITIONS: usize = 4096; + +/// IVF partitions for a table of `rows`: about one per sqrt(rows), at least one. +pub fn partitions_for(rows: usize) -> usize { + ((rows as f64).sqrt().round() as usize).clamp(1, MAX_PARTITIONS) +} + +/// The metric the table was created with, from the vector field's +/// `distance_metric` metadata as the S3 Vectors gateway writes it; L2 when +/// absent. +pub fn distance_metric(dataset: &Dataset, vector_column: &str) -> DistanceType { + match dataset + .schema() + .field(vector_column) + .and_then(|f| f.metadata.get("distance_metric").map(String::as_str)) + { + Some("COSINE") => DistanceType::Cosine, + Some("DOT_PRODUCT") => DistanceType::Dot, + _ => DistanceType::L2, + } +} + +/// Names of the indices built on `vector_column`, each once. +pub async fn index_names(dataset: &Dataset, vector_column: &str) -> Result> { + let Some(field) = dataset.schema().field(vector_column) else { + return Ok(Vec::new()); + }; + let mut names: Vec = dataset + .load_indices() + .await? + .iter() + .filter(|i| i.fields == [field.id]) + .map(|i| i.name.clone()) + .collect(); + names.sort(); + names.dedup(); + Ok(names) +} + +/// Build an IVF_FLAT index named `name` over `vector_column` in the metric the +/// table was created with. `replace` rebuilds an index of that name; the +/// dataset is left at the new version. +pub async fn build( + dataset: &mut Dataset, + vector_column: &str, + name: &str, + partitions: usize, + replace: bool, +) -> Result<()> { + let params = VectorIndexParams::ivf_flat(partitions, distance_metric(dataset, vector_column)); + dataset + .create_index( + &[vector_column], + IndexType::Vector, + Some(name.to_string()), + ¶ms, + replace, + ) + .await + .with_context(|| format!("build index {name} with {partitions} partitions"))?; + dataset.checkout_latest().await?; + Ok(()) +} + +/// Extend the indices on `vector_column` over rows written since they were +/// built: a new delta segment, or with `retrain` one index trained again on +/// all rows. Indices on other columns are left alone. +pub async fn optimize(dataset: &mut Dataset, vector_column: &str, retrain: bool) -> Result<()> { + let mut options = if retrain { + OptimizeOptions::retrain() + } else { + OptimizeOptions::append() + }; + options.index_names = Some(index_names(dataset, vector_column).await?); + dataset.optimize_indices(&options).await?; + dataset.checkout_latest().await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn partitions_follow_the_square_root() { + assert_eq!(partitions_for(0), 1); + assert_eq!(partitions_for(64), 8); + assert_eq!(partitions_for(50_000), 224); + assert_eq!(partitions_for(1_000_000), 1000); + assert_eq!(partitions_for(100_000_000), MAX_PARTITIONS); + } +} diff --git a/test/s3vectors/Makefile b/test/s3vectors/Makefile index 5a390bb950..77048ada54 100644 --- a/test/s3vectors/Makefile +++ b/test/s3vectors/Makefile @@ -1,5 +1,6 @@ WEED := ./weed VEC := ./weed-vector +CLI := ./weed-vector-cli WORKER_MANIFEST := ../../rust/sw-vector-distributed/Cargo.toml ifneq ($(wildcard $(WORKER_MANIFEST)),) WORKER := ./weed-vector-worker @@ -22,6 +23,10 @@ $(VEC): FORCE cd ../../seaweed-vector && cargo build --locked --profile $(CARGO_PROFILE) -p weed-vector-server cp $(CARGO_OUTPUT)/weed-vector $(VEC) +$(CLI): FORCE + cd ../../seaweed-vector && cargo build --locked --profile $(CARGO_PROFILE) -p weed-vector-core --bin weed-vector-cli + cp $(CARGO_OUTPUT)/weed-vector-cli $(CLI) + ifneq ($(WORKER),) $(WORKER): $(VEC) FORCE cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../rust/Cargo.toml -p sw-vector-distributed --bin weed-vector @@ -39,8 +44,8 @@ build: $(WEED) $(VEC) $(WORKER) $(VOLUME) test: build ./run_test.sh -perf: build +perf: build $(CLI) ./run_perf.sh clean: - rm -f $(WEED) $(VEC) ./weed-vector-worker $(VOLUME) + rm -f $(WEED) $(VEC) $(CLI) ./weed-vector-worker $(VOLUME) diff --git a/test/s3vectors/conftest.py b/test/s3vectors/conftest.py index 94a2fac260..e54816b8a4 100644 --- a/test/s3vectors/conftest.py +++ b/test/s3vectors/conftest.py @@ -330,7 +330,8 @@ def run_cluster(param, tmp_path_factory, identities): processes, "gateway did not discover both search workers", ) - yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory} + yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory, + "catalog": f"http://127.0.0.1:{lance}"} @pytest.fixture(scope="module") diff --git a/test/s3vectors/perf_test.py b/test/s3vectors/perf_test.py index 94fa4eb813..5b7652c4db 100644 --- a/test/s3vectors/perf_test.py +++ b/test/s3vectors/perf_test.py @@ -12,10 +12,19 @@ a larger profile: PERF_CONCURRENCY client threads for the throughput scenario ("1,4,8") PERF_RECALL queries checked against an exact search (20) PERF_REPORT JSON report path (perf-results.json in the cluster dir) + PERF_INDEX 1 builds the vector index after ingest with weed-vector-cli, + so the queries run through it (PERF_PARTITIONS overrides + the partition count) + PERF_CLUSTERS draw vectors around this many shared centres instead of + independently (0); PERF_NOISE is the spread around a + centre (0.3). Independent Gaussian vectors have no + locality for an IVF index to use, so recall through one + only tracks the fraction of partitions probed. Each timed scenario starts with one untimed call. Timings only report; the assertions check correctness (recall against an exact client-side search, -complete listings), never speed. +complete listings), never speed. Recall through an index is approximate, so +PERF_INDEX runs report it without asserting it. """ import concurrent.futures @@ -27,6 +36,7 @@ import platform import random import re import statistics +import subprocess import time import uuid @@ -34,6 +44,7 @@ import boto3 from botocore.config import Config import pytest +HERE = Path(__file__).resolve().parent VECTORS = int(os.environ.get("PERF_VECTORS", "10000")) DIM = int(os.environ.get("PERF_DIM", "128")) BATCH = int(os.environ.get("PERF_BATCH", "500")) @@ -42,6 +53,9 @@ TOPK = int(os.environ.get("PERF_TOPK", "10")) CONCURRENCY = [int(c) for c in os.environ.get("PERF_CONCURRENCY", "1,4,8").split(",")] CATEGORIES = 10 RECALL_SAMPLE = int(os.environ.get("PERF_RECALL", "20")) +INDEXED = bool(os.environ.get("PERF_INDEX")) +CLUSTERS = int(os.environ.get("PERF_CLUSTERS", "0")) +NOISE = float(os.environ.get("PERF_NOISE", "0.3")) INDEX = "perf" REPORT = {} @@ -70,16 +84,27 @@ def client(endpoint): max_pool_connections=max(CONCURRENCY) + 2)) -def exact_top_k(data, query, k, keep=None): - """Cosine top-k over the ingested vectors, the reference for recall.""" +def recall(data, query, k, hits, keep=None): + """Hits in the exact cosine top-k over the ingested vectors. Hits within + 1e-6 of the k-th best score fill the slots the top-k leaves for that tie: + tied vectors are interchangeable, and the server ranks in float32.""" qn = math.sqrt(sum(x * x for x in query)) - scored = [] - for i, (vector, norm) in enumerate(data): - if keep is not None and not keep(i): - continue - scored.append((sum(a * b for a, b in zip(vector, query)) / (norm * qn), i)) - scored.sort(reverse=True) - return {f"v{i}" for _, i in scored[:k]} + + def score(i): + vector, norm = data[i] + return sum(a * b for a, b in zip(vector, query)) / (norm * qn) + + def eligible(i): + return keep is None or keep(i) + + best = sorted((score(i) for i in range(len(data)) if eligible(i)), reverse=True)[:k] + found = {int(h["key"][1:]) for h in hits} + cutoff = best[-1] + scores = [score(i) for i in found if eligible(i)] + above = sum(1 for s in scores if s > cutoff + 1e-6) + tied = sum(1 for s in scores if abs(s - cutoff) <= 1e-6) + slots = len(best) - sum(1 for s in best if s > cutoff + 1e-6) + return (above + min(tied, slots)) / len(best) def list_all(s3v, args, **extra): @@ -96,19 +121,23 @@ def assert_complete(keys): assert len(keys) == VECTORS and set(keys) == {f"v{i}" for i in range(VECTORS)} -def recall(expected, hits): - return len(expected & {h["key"] for h in hits}) / len(expected) - - @pytest.fixture(scope="module") def perf(cluster): assert 1 <= BATCH <= 500 and VECTORS >= BATCH rng = random.Random(1) + centres = [[rng.gauss(0, 1) for _ in range(DIM)] for _ in range(CLUSTERS)] + + def vector(): + if not centres: + return [round(rng.gauss(0, 1), 5) for _ in range(DIM)] + centre = rng.choice(centres) + return [round(c + NOISE * rng.gauss(0, 1), 5) for c in centre] + data = [] for _ in range(VECTORS): - vector = [round(rng.gauss(0, 1), 5) for _ in range(DIM)] - data.append((vector, math.sqrt(sum(x * x for x in vector)))) - queries = [[round(rng.gauss(0, 1), 5) for _ in range(DIM)] for _ in range(QUERIES)] + v = vector() + data.append((v, math.sqrt(sum(x * x for x in v)))) + queries = [vector() for _ in range(QUERIES)] s3v = client(cluster["endpoint"]) bucket = "perf-" + uuid.uuid4().hex[:12] @@ -118,10 +147,11 @@ def perf(cluster): args = {"vectorBucketName": bucket, "indexName": INDEX} combo = re.sub(r"^vectors-|\d+$", "", cluster["directory"].name) REPORT[combo] = {"config": {"vectors": VECTORS, "dim": DIM, "batch": BATCH, "queries": QUERIES, + "clusters": CLUSTERS, "noise": NOISE, "top_k": TOPK, "concurrency": CONCURRENCY, "cpus": os.cpu_count(), "load_1m": os.getloadavg()[0], "machine": platform.machine()}} yield {"client": s3v, "endpoint": cluster["endpoint"], "args": args, "data": data, - "queries": queries, "report": REPORT[combo]} + "queries": queries, "report": REPORT[combo], "cluster": cluster, "bucket": bucket} s3v.delete_index(**args) s3v.delete_vector_bucket(vectorBucketName=bucket) path = Path(os.environ.get("PERF_REPORT", cluster["directory"] / "perf-results.json")) @@ -158,6 +188,23 @@ def test_upsert_existing(perf): assert_complete(list_all(s3v, args)) +def test_build_index(perf): + if not INDEXED: + pytest.skip("PERF_INDEX not set") + cli = Path(os.environ.get("VECTOR_CLI_BINARY", HERE / "weed-vector-cli")).resolve() + assert cli.is_file(), f"missing {cli}; run make ./weed-vector-cli" + command = [str(cli), "--namespace", perf["cluster"]["catalog"], "--access-key", "test", + "--secret-key", "test", "index", "build", f"{perf['bucket']}$vector${INDEX}"] + if os.environ.get("PERF_PARTITIONS"): + command += ["--partitions", os.environ["PERF_PARTITIONS"]] + seconds, out = timed(lambda: subprocess.run(command, check=True, capture_output=True, text=True)) + report = json.loads(out.stdout) + coverage = report["coverage"] + assert len(coverage) == 1 and coverage[0]["unindexed_rows"] == 0, coverage + perf["report"]["index"] = {"seconds": seconds, "segments": coverage[0]["segments"], + "partitions": int(os.environ.get("PERF_PARTITIONS", "0")) or None} + + def run_queries(perf, name, **extra): s3v, args, data, queries = perf["client"], perf["args"], perf["data"], perf["queries"] latencies, recalls = [], [] @@ -169,10 +216,12 @@ def run_queries(perf, name, **extra): seconds, response = timed(lambda: s3v.query_vectors( **args, queryVector={"float32": query}, topK=TOPK, returnDistance=True, **extra)) latencies.append(seconds) + if keep: + assert all(keep(int(h["key"][1:])) for h in response["vectors"]), response["vectors"] if n < RECALL_SAMPLE: - recalls.append(recall(exact_top_k(data, query, TOPK, keep), response["vectors"])) + recalls.append(recall(data, query, TOPK, response["vectors"], keep)) perf["report"][name] = {**summarize(latencies), "recall_at_k": statistics.fmean(recalls) if recalls else None} - assert not recalls or statistics.fmean(recalls) >= 0.9 + assert INDEXED or not recalls or statistics.fmean(recalls) >= 0.9 def test_query(perf): @@ -258,6 +307,9 @@ def print_report(combo, report, path): cfg = report["config"] lines = [f"\n== S3 Vectors perf [{combo}]: {cfg['vectors']} x {cfg['dim']}d cosine, " f"topK {cfg['top_k']}, {cfg['cpus']} CPUs, load {cfg['load_1m']:.1f}"] + index = report.get("index") + if index: + lines.append(f"index built in {index['seconds']:.1f} s, {index['segments']} segment(s)") ingest = report.get("ingest") if ingest: lines.append(f"ingest {ingest['vectors_per_s']:9.0f} vectors/s "