diff --git a/seaweed-vector/crates/core/src/index.rs b/seaweed-vector/crates/core/src/index.rs index ca19d0c8b9..9d820f8c87 100644 --- a/seaweed-vector/crates/core/src/index.rs +++ b/seaweed-vector/crates/core/src/index.rs @@ -1,18 +1,60 @@ //! Sizing the vector index an S3 Vectors table searches through, and the //! probes a query spends on it. Building lives in the worker crate. +use anyhow::Result; +use lance::Dataset; +use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; + 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)) +/// How many partitions a query searches: `min` always, and under a prefilter +/// up to `max` until it has k rows. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Probes { + pub min: usize, + pub max: usize, +} + +/// Probes for an index of `partitions`: a tenth of them, at least +/// [`MIN_PROBES`], never more than exist; a filtered query may widen to four +/// times that. +pub fn probes_for(partitions: usize) -> Probes { + let partitions = partitions.max(1); + let min = MIN_PROBES.max(partitions.div_ceil(10)).min(partitions); + Probes { + min, + max: (min * 4).min(partitions), + } +} + +/// The probes a query of `vector_column` uses when it names none, or `None` +/// without an index. The count comes from the most partitions any index on +/// the column has, whichever of them the scan ends up searching. +/// It is read from the opened index, which the scan then finds loaded; +/// Lance's index statistics cost several milliseconds a query. +pub async fn default_probes(dataset: &Dataset, vector_column: &str) -> Result> { + let Some(field) = dataset.schema().field(vector_column) else { + return Ok(None); + }; + let mut partitions = None; + for index in dataset.load_indices().await?.iter() { + if index.fields != [field.id] { + continue; + } + let opened = dataset + .open_vector_index( + vector_column, + &index.uuid, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await?; + partitions = Some(opened.total_partitions().max(partitions.unwrap_or(0))); + } + Ok(partitions.map(probes_for)) } #[cfg(test)] @@ -29,11 +71,17 @@ mod tests { #[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); + assert_eq!(probes_for(1), Probes { min: 1, max: 1 }); + assert_eq!(probes_for(8), Probes { min: 8, max: 8 }); + assert_eq!(probes_for(64), Probes { min: 8, max: 32 }); + assert_eq!(probes_for(224), Probes { min: 23, max: 92 }); + assert_eq!( + probes_for(4096), + Probes { + min: 410, + max: 1640 + } + ); } fn exact(vector: &[f32], k: usize) -> Query { @@ -73,6 +121,107 @@ mod tests { } } + /// Pseudo-random rows, so partitions are not degenerate; a different seed + /// gives vectors the table does not hold. + fn scattered(rows: usize, dim: usize, seed: u64) -> Vec> { + let mut state: u64 = seed; + (0..rows) + .map(|_| { + (0..dim) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 40) as f32) / (1u64 << 24) as f32 + }) + .collect() + }) + .collect() + } + + #[tokio::test] + async fn default_probes_follow_the_index_and_find_the_neighbours() { + let dir = TempDir::new().unwrap(); + let rows = scattered(4000, 8, 0x9E3779B97F4A7C15); + let tags: Vec<&str> = (0..4000) + .map(|i| if i % 100 == 0 { "rare" } else { "common" }) + .collect(); + let mut dataset = testing::write_dataset(&dir, &rows, &tags).await; + let table = VectorTable::resolve(&dataset, None, None).unwrap(); + assert_eq!(default_probes(&dataset, "vector").await.unwrap(), None); + build(&mut dataset, &table.vector_column, INDEX_NAME, 64, false) + .await + .unwrap(); + assert_eq!( + default_probes(&dataset, "vector").await.unwrap(), + Some(probes_for(64)) + ); + + let mut found = 0; + let total = 20 * 10; + for vector in &scattered(20, 8, 0xD1B54A32D192ED03) { + let q = exact(vector, 10); + let flat: std::collections::HashSet = query::flat(&dataset, &table, &q) + .await + .unwrap() + .iter() + .map(|h| h.row_id) + .collect(); + let ann = query::ann(&dataset, &table, &q).await.unwrap(); + assert_eq!(ann.len(), 10); + found += ann.iter().filter(|h| flat.contains(&h.row_id)).count(); + } + assert!(found * 10 >= total * 9, "recall {found}/{total}"); + + // Eight probes do not hold ten of the 1% rare rows; a filtered query + // widens past the minimum until it has them. + let filtered = query::ann( + &dataset, + &table, + &Query { + filter: Some("tag = 'rare'".to_string()), + ..exact(&rows[0], 10) + }, + ) + .await + .unwrap(); + assert_eq!(filtered.len(), 10); + } + + #[tokio::test] + async fn default_probes_take_the_largest_index() { + use lance::index::DatasetIndexExt; + use lance::index::vector::VectorIndexParams; + use lance_index::IndexType; + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &scattered(1000, 8, 0x9E3779B97F4A7C15), &["a"; 1000]) + .await; + testing::append_dataset( + &mut dataset, + &scattered(1000, 8, 0xD1B54A32D192ED03), + &["b"; 1000], + ) + .await; + for (name, fragment, partitions) in [("a-small", 0u32, 2), ("z-large", 1u32, 100)] { + dataset + .create_index_builder( + &["vector"], + IndexType::Vector, + &VectorIndexParams::ivf_flat(partitions, DistanceType::L2), + ) + .name(name.to_string()) + .fragments(vec![fragment]) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + } + assert_eq!( + default_probes(&dataset, "vector").await.unwrap(), + Some(probes_for(100)) + ); + } + #[tokio::test] async fn build_covers_every_fragment() { let dir = TempDir::new().unwrap(); diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs index 9cf7ef9ee7..c8e671851b 100644 --- a/seaweed-vector/crates/core/src/query.rs +++ b/seaweed-vector/crates/core/src/query.rs @@ -120,8 +120,28 @@ async fn scanner( crate::filter::apply(&mut scanner, filter)?; scanner.prefilter(!query.postfilter); } - if let Some(nprobes) = query.nprobes { - scanner.nprobes(nprobes); + // An explicitly empty segment list searches no index, so there is + // nothing to size probes for. + let no_segments = query.index_segments.as_ref().is_some_and(|s| s.is_empty()); + if use_index && !no_segments { + let probes = match query.nprobes { + Some(n) => Some(crate::index::Probes { min: n, max: n }), + None => crate::index::default_probes(dataset, &table.vector_column).await?, + }; + if let Some(probes) = probes { + // Lance searches `min` partitions, and under a prefilter keeps + // going up to `max` until it has k rows. An open bound would + // make it prepare every partition for that, which costs as much + // as the exact scan it is meant to avoid. + scanner.minimum_nprobes(probes.min); + scanner.maximum_nprobes(probes.max); + tracing::debug!( + min = probes.min, + max = probes.max, + filtered = query.filter.is_some(), + "ann scan" + ); + } } let segments = query.index_segments.as_deref().unwrap_or(&[]); if use_index && !segments.is_empty() { diff --git a/seaweed-vector/crates/server/src/main.rs b/seaweed-vector/crates/server/src/main.rs index 94f89563ad..d9eeb21267 100644 --- a/seaweed-vector/crates/server/src/main.rs +++ b/seaweed-vector/crates/server/src/main.rs @@ -36,6 +36,10 @@ struct Args { /// Memory for caching reads of write-once table files, in MiB; 0 disables. #[arg(long, env = "WEED_VECTOR_READ_CACHE_MB", default_value_t = 512)] read_cache_mb: u64, + /// Partitions each query probes in a table's vector index. Default: a + /// tenth of the index's partitions, at least 8; more under a filter. + #[arg(long, env = "WEED_VECTOR_NPROBES")] + nprobes: Option, /// The `jwt.filer_signing.key` of the cluster. The Lance namespace vends /// credentials that write index data only to a caller holding it. #[arg(long, env = "WEED_FILER_SIGNING_KEY")] @@ -154,7 +158,8 @@ async fn main() -> Result<()> { let client = NamespaceClient::new(args.namespace.clone()) .with_fallback(&fallback) .with_signing_key(signing_key.as_ref()); - let mut service = VectorServer::new(client, fallback); + let mut service = VectorServer::new(client, fallback) + .with_nprobes(args.nprobes.map(std::num::NonZeroUsize::get)); if args.read_cache_mb > 0 { service = service.with_read_cache(ReadCache::new(args.read_cache_mb << 20)); } @@ -192,3 +197,15 @@ async fn main() -> Result<()> { } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nprobes_must_be_positive() { + assert!(Args::try_parse_from(["weed-vector", "--nprobes", "0"]).is_err()); + let args = Args::try_parse_from(["weed-vector", "--nprobes", "12"]).unwrap(); + assert_eq!(args.nprobes.map(std::num::NonZeroUsize::get), Some(12)); + } +} diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index 66d2c44d9e..6237b66b44 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -374,6 +374,7 @@ pub struct VectorServer { client: NamespaceClient, fallback: FallbackOptions, read_cache: Option>, + nprobes: Option, } impl VectorServer { @@ -382,9 +383,16 @@ impl VectorServer { client, fallback, read_cache: None, + nprobes: None, } } + /// Partitions every query probes, instead of the index's own default. + pub fn with_nprobes(mut self, nprobes: Option) -> Self { + self.nprobes = nprobes; + self + } + /// Serve repeated reads of write-once table files from `cache`. pub fn with_read_cache(mut self, cache: ReadCache) -> Self { self.read_cache = Some(Arc::new(cache)); @@ -669,7 +677,7 @@ impl VectorService for VectorServer { }, postfilter: req.postfilter, distance, - nprobes: None, + nprobes: self.nprobes, index_segments: None, fragments: None, };