diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index 2d13f80da6..8590908110 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6597,6 +6597,7 @@ dependencies = [ "lance-table", "prometheus", "reqwest 0.12.28", + "roaring", "seaweed-common", "seaweed-worker-core", "seaweed-worker-sort", diff --git a/seaweed-vector/crates/core/src/coverage.rs b/seaweed-vector/crates/core/src/coverage.rs index 7849215497..42a58499cd 100644 --- a/seaweed-vector/crates/core/src/coverage.rs +++ b/seaweed-vector/crates/core/src/coverage.rs @@ -13,6 +13,8 @@ pub struct IndexCoverage { pub index_name: String, /// Physical index segments: the IndexMetadata deltas sharing the name. pub segments: usize, + /// IVF partitions of the largest segment. + pub partitions: usize, /// False when any delta of this index lacks a fragment bitmap: the index /// exists but which fragments it covers is unknown. pub coverage_known: bool, @@ -28,6 +30,7 @@ pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result>> = BTreeMap::new(); + let mut partitions: BTreeMap = BTreeMap::new(); for index in dataset.load_indices().await?.iter() { if index.fields != [field.id] { continue; @@ -36,6 +39,9 @@ pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result Result Probes { } } +/// Partitions of one index segment, read from the opened index. Lance's index +/// statistics cost several milliseconds and may rewrite an old manifest. +pub async fn segment_partitions( + dataset: &Dataset, + vector_column: &str, + segment: &uuid::Uuid, +) -> Result { + let opened = dataset + .open_vector_index( + vector_column, + segment, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await?; + Ok(opened.total_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. +/// the column has, whichever of them the scan ends up searching; the scan +/// then finds the index it was read from already open. pub async fn default_probes(dataset: &Dataset, vector_column: &str) -> Result> { let Some(field) = dataset.schema().field(vector_column) else { return Ok(None); @@ -45,14 +61,8 @@ pub async fn default_probes(dataset: &Dataset, vector_column: &str) -> Result(), + 3, + "{:?}", + cover.deltas + ); + let per_name = coverage(&dataset, "vector").await.unwrap(); + let untouched = per_name + .iter() + .filter(|c| c.segments == 1 && c.indexed_fragments.len() == 1) + .count(); + assert_eq!(untouched, 1, "{per_name:?}"); + for vector in testing::vectors(3, DIM) { + assert_same_top_k(&dataset, &vector, 10).await; + } + } + + /// A retrain is due for one IVF_FLAT name the table outgrew, never for + /// shards sized each for their own fragments or for another index type, + /// and an explicit retrain of shards leaves them as they are. + #[tokio::test] + async fn retrain_is_due_for_one_ivf_flat_name_only() { + use lance::index::DatasetIndexExt; + use lance::index::vector::VectorIndexParams; + use lance_index::IndexType; + use weed_lance_worker::vector_index::{coverage as covered, retrain_due}; + let grown = |dir: &TempDir| { + let path = dir.path().to_str().unwrap().to_string(); + async move { + let mut dataset = + testing::write_dataset_at(&path, &testing::vectors(300, DIM), &["a"; 300]) + .await; + testing::append_dataset(&mut dataset, &testing::vectors(300, DIM), &["b"; 300]) + .await; + dataset + } + }; + + let dir = TempDir::new().unwrap(); + let mut dataset = grown(&dir).await; + build(&mut dataset, "vector", INDEX_NAME, 2, false) + .await + .unwrap(); + assert!(retrain_due(&dataset, "vector").await.unwrap()); + + let dir = TempDir::new().unwrap(); + let mut dataset = grown(&dir).await; + let params = VectorIndexParams::ivf_pq(2, 8, 2, DistanceType::L2, 10); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + assert!(!retrain_due(&dataset, "vector").await.unwrap()); + + let dir = TempDir::new().unwrap(); + let mut dataset = grown(&dir).await; + let params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + for (name, fragment) in [("shard-0", 0u32), ("shard-1", 1u32)] { + dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name(name.to_string()) + .fragments(vec![fragment]) + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + } + assert!(!retrain_due(&dataset, "vector").await.unwrap()); + let before = covered(&dataset, "vector").await.unwrap(); + assert!(optimize(&mut dataset, "vector", true).await.is_err()); + assert_eq!(covered(&dataset, "vector").await.unwrap(), before); + let per_name = coverage(&dataset, "vector").await.unwrap(); + assert!( + per_name.iter().all(|c| c.indexed_fragments.len() == 1), + "{per_name:?}" + ); + } + #[tokio::test] async fn compaction_keeps_the_index_covering() { let dir = TempDir::new().unwrap(); diff --git a/seaweed-worker/Cargo.lock b/seaweed-worker/Cargo.lock index e1a57e6611..f4c08a3c79 100644 --- a/seaweed-worker/Cargo.lock +++ b/seaweed-worker/Cargo.lock @@ -6518,6 +6518,7 @@ dependencies = [ "lance-table", "prometheus", "reqwest 0.12.28", + "roaring", "seaweed-common", "seaweed-worker-core", "seaweed-worker-sort", diff --git a/seaweed-worker/crates/lance/Cargo.toml b/seaweed-worker/crates/lance/Cargo.toml index 0ce5f4c9cd..951200d445 100644 --- a/seaweed-worker/crates/lance/Cargo.toml +++ b/seaweed-worker/crates/lance/Cargo.toml @@ -28,6 +28,7 @@ prometheus.workspace = true lance = { version = "10", default-features = false, features = ["aws"] } lance-index = "10" lance-linalg = "10" +roaring = "0.11" # 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/jobs/indices.rs b/seaweed-worker/crates/lance/src/jobs/indices.rs index 24f87b0c68..87ce12d723 100644 --- a/seaweed-worker/crates/lance/src/jobs/indices.rs +++ b/seaweed-worker/crates/lance/src/jobs/indices.rs @@ -6,8 +6,8 @@ use lance::index::DatasetIndexExt; use lance_index::optimize::OptimizeOptions; use seaweed_worker_core::config_form::{form, int_or, int_value, number_field}; use seaweed_worker_core::pb::{ - ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, JobCompleted, - JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor, + AdminRuntimeDefaults, ConfigValue, DetectionComplete, DetectionProposals, ExecuteJobRequest, + JobCompleted, JobProgressUpdate, JobProposal, JobResult, JobTypeCapability, JobTypeDescriptor, RunDetectionRequest, }; use seaweed_worker_core::{DetectionSender, ExecutionSender, JobHandler}; @@ -15,9 +15,13 @@ use tracing::warn; use crate::catalog::{NamespaceClient, SigningKey, parse_id}; use crate::dataset::{self, OpenTable}; +use crate::jobs::vector_index::is_vector_table; use crate::jobs::{clamp, string_list, table_id}; +use crate::vector_index; pub const JOB_TYPE: &str = "lance_optimize_indices"; +/// The vector column the S3 Vectors API writes. +const VECTOR_COLUMN: &str = "vector"; const DEFAULT_MAX_UNINDEXED_ROWS: i64 = 100_000; @@ -71,16 +75,34 @@ impl OptimizeIndicesHandler { } } -/// Rows no index covers, summed across a table's indices. A table with no -/// indices at all has nothing to optimize, which is different from a table whose -/// indices have fallen behind. -async fn unindexed_rows(table: &OpenTable) -> Result> { +/// Rows the table's indices cover and rows they miss. A vector table's indices +/// are taken together, since a sharded build gives one name per shard; any +/// other table reports the index that misses the most. A table with no +/// indices at all has nothing to optimize, which is different from a table +/// whose indices have fallen behind. +struct IndexRows { + indexed: u64, + unindexed: u64, +} + +async fn index_rows(table: &OpenTable, vector_table: bool) -> Result> { + if vector_table { + return Ok(vector_index::coverage(&table.dataset, VECTOR_COLUMN) + .await? + .map(|cover| IndexRows { + indexed: cover.indexed_rows, + unindexed: cover.unindexed_rows, + })); + } let indices = table.dataset.load_indices().await?; if indices.is_empty() { return Ok(None); } - let mut worst = 0u64; + let mut worst = IndexRows { + indexed: 0, + unindexed: 0, + }; let mut names: Vec = indices.iter().map(|index| index.name.clone()).collect(); names.sort(); names.dedup(); @@ -88,15 +110,27 @@ async fn unindexed_rows(table: &OpenTable) -> Result> { let raw = table.dataset.index_statistics(&name).await?; let stats: serde_json::Value = serde_json::from_str(&raw) .with_context(|| format!("parse index statistics for {name}"))?; - let unindexed = stats - .get("num_unindexed_rows") - .and_then(|value| value.as_u64()) - .unwrap_or(0); - worst = worst.max(unindexed); + let count = |key: &str| stats.get(key).and_then(|value| value.as_u64()).unwrap_or(0); + if count("num_unindexed_rows") >= worst.unindexed { + worst = IndexRows { + indexed: count("num_indexed_rows"), + unindexed: count("num_unindexed_rows"), + }; + } } Ok(Some(worst)) } +/// A vector index also reindexes once a tenth of it is uncovered, so a small +/// index does not wait for the absolute budget. +fn effective_budget(budget: u64, rows: &IndexRows, vector_table: bool) -> u64 { + if vector_table { + budget.min(rows.indexed / 10) + } else { + budget + } +} + #[async_trait] impl JobHandler for OptimizeIndicesHandler { fn capability(&self) -> JobTypeCapability { @@ -138,6 +172,12 @@ impl JobHandler for OptimizeIndicesHandler { defaults.clone(), )), worker_default_values: defaults, + // A retrain is a whole build, so it gets a build's time. + admin_runtime_defaults: Some(AdminRuntimeDefaults { + job_type_max_runtime_seconds: super::vector_index::EXECUTION_TIMEOUT_SECONDS, + execution_timeout_seconds: super::vector_index::EXECUTION_TIMEOUT_SECONDS, + ..Default::default() + }), ..Default::default() } } @@ -175,8 +215,9 @@ impl JobHandler for OptimizeIndicesHandler { if let Some(counters) = &self.metrics { counters.worker.object_seen(JOB_TYPE); } - let unindexed = match unindexed_rows(&table).await { - Ok(Some(unindexed)) => unindexed, + let vector_table = is_vector_table(&id); + let rows = match index_rows(&table, vector_table).await { + Ok(Some(rows)) => rows, Ok(None) => continue, Err(err) => { if let Some(counters) = &self.metrics { @@ -186,11 +227,17 @@ impl JobHandler for OptimizeIndicesHandler { continue; } }; + let budget = effective_budget(budget, &rows, vector_table); + let unindexed = rows.unindexed; if unindexed <= budget { continue; } let mut parameters: HashMap = HashMap::new(); parameters.insert("table_id".to_string(), string_list(&id)); + if vector_table { + // The run may turn out to be a retrain. + super::vector_index::declare_build_runtime(&mut parameters); + } proposals.push(JobProposal { proposal_id: format!("{JOB_TYPE}:{encoded}"), dedupe_key: format!("{JOB_TYPE}:{encoded}"), @@ -238,27 +285,46 @@ impl JobHandler for OptimizeIndicesHandler { let client = self.client(); let mut table = dataset::open(&client, &id, &self.fallback).await?; - let before = unindexed_rows(&table).await?.unwrap_or(0); + let vector_table = is_vector_table(&id); + let before = index_rows(&table, vector_table) + .await? + .map_or(0, |rows| rows.unindexed); + // A vector index whose table has outgrown its partitions is built + // again with partitions for the rows it now holds. + let retrain = + vector_table && vector_index::retrain_due(&table.dataset, VECTOR_COLUMN).await?; sender.send_progress(JobProgressUpdate { request_id: request.request_id.clone(), job_id: job.job_id.clone(), job_type: JOB_TYPE.to_string(), progress_percent: 10.0, - stage: format!("indexing {before} uncovered rows"), + stage: if retrain { + "retraining the vector index".to_string() + } else { + format!("indexing {before} uncovered rows") + }, ..Default::default() })?; - // Merging every delta keeps read latency from drifting as the table is - // reindexed again and again; leaving them unmerged is how an index ends - // up fast to write and slow to search. - table - .dataset - .optimize_indices(&OptimizeOptions::default()) - .await - .with_context(|| format!("optimize indices of {}", table.location))?; + if vector_table { + vector_index::optimize(&mut table.dataset, VECTOR_COLUMN, retrain) + .await + .with_context(|| format!("optimize the vector index of {}", table.location))?; + } else { + // Merging every delta keeps read latency from drifting as the table + // is reindexed again and again; leaving them unmerged is how an + // index ends up fast to write and slow to search. + table + .dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .with_context(|| format!("optimize indices of {}", table.location))?; + } - let after = unindexed_rows(&table).await?.unwrap_or(0); + let after = index_rows(&table, vector_table) + .await? + .map_or(0, |rows| rows.unindexed); if let Some(counters) = &self.metrics { counters.rows_indexed.inc_by(before.saturating_sub(after)); } @@ -268,6 +334,10 @@ impl JobHandler for OptimizeIndicesHandler { int_value(before as i64), ); output.insert("unindexed_rows_after".to_string(), int_value(after as i64)); + output.insert( + "retrained".to_string(), + seaweed_worker_core::config_form::bool_value(retrain), + ); sender.send_completed(JobCompleted { request_id: request.request_id.clone(), @@ -284,3 +354,32 @@ impl JobHandler for OptimizeIndicesHandler { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_retrain_gets_a_builds_execution_timeout() { + let descriptor = OptimizeIndicesHandler::new(String::new()).descriptor(); + let defaults = descriptor.admin_runtime_defaults.unwrap(); + assert_eq!(defaults.execution_timeout_seconds, 3600); + assert_eq!(defaults.job_type_max_runtime_seconds, 3600); + assert!(!defaults.enabled); + } + + #[test] + fn vector_tables_reindex_at_a_tenth_of_their_rows() { + let small = IndexRows { + indexed: 60_000, + unindexed: 0, + }; + assert_eq!(effective_budget(100_000, &small, true), 6_000); + assert_eq!(effective_budget(100_000, &small, false), 100_000); + let large = IndexRows { + indexed: 5_000_000, + unindexed: 0, + }; + assert_eq!(effective_budget(100_000, &large, true), 100_000); + } +} diff --git a/seaweed-worker/crates/lance/src/jobs/vector_index.rs b/seaweed-worker/crates/lance/src/jobs/vector_index.rs index f07b43d084..9941907e2d 100644 --- a/seaweed-worker/crates/lance/src/jobs/vector_index.rs +++ b/seaweed-worker/crates/lance/src/jobs/vector_index.rs @@ -27,7 +27,17 @@ const VECTOR_NAMESPACE: &str = "vector"; const VECTOR_COLUMN: &str = "vector"; /// A build trains on the whole table, so the admin's 90 seconds are not enough. /// The run window is as long, or it would cut an attempt short first. -const EXECUTION_TIMEOUT_SECONDS: i32 = 3600; +pub(crate) const EXECUTION_TIMEOUT_SECONDS: i32 = 3600; + +/// Declares a build's time on a proposal. A run started by hand while the +/// schedule is off has no policy to take the timeout from, and the admin +/// gives a job at least the runtime its proposal declares. +pub(crate) fn declare_build_runtime(parameters: &mut HashMap) { + parameters.insert( + "estimated_runtime_seconds".to_string(), + int_value(EXECUTION_TIMEOUT_SECONDS as i64), + ); +} /// Builds the vector index of an S3 Vectors table once it holds enough rows /// for an exact scan to be the slow way. Tables below the threshold, and rows @@ -215,6 +225,7 @@ impl JobHandler for VectorIndexHandler { } let mut parameters: HashMap = HashMap::new(); parameters.insert("table_id".to_string(), string_list(&id)); + declare_build_runtime(&mut parameters); proposals.push(JobProposal { proposal_id: format!("{JOB_TYPE}:{encoded}"), dedupe_key: format!("{JOB_TYPE}:{encoded}"), @@ -343,6 +354,16 @@ mod tests { assert_eq!(skip_reason(false, 60_000), None); } + #[test] + fn a_proposal_declares_a_builds_runtime() { + let mut parameters = HashMap::new(); + declare_build_runtime(&mut parameters); + assert_eq!( + seaweed_worker_core::config_form::int_or(¶meters, "estimated_runtime_seconds", 0), + 3600 + ); + } + #[test] fn a_build_gets_longer_than_the_default_execution_timeout() { let descriptor = VectorIndexHandler::new(String::new()).descriptor(); diff --git a/seaweed-worker/crates/lance/src/vector_index.rs b/seaweed-worker/crates/lance/src/vector_index.rs index f3356273ae..c2b43fb9b0 100644 --- a/seaweed-worker/crates/lance/src/vector_index.rs +++ b/seaweed-worker/crates/lance/src/vector_index.rs @@ -1,6 +1,8 @@ //! 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 std::collections::BTreeMap; + use anyhow::{Context, Result}; use lance::Dataset; use lance::index::DatasetIndexExt; @@ -8,10 +10,14 @@ use lance::index::vector::VectorIndexParams; use lance_index::IndexType; use lance_index::optimize::OptimizeOptions; use lance_linalg::distance::DistanceType; +use roaring::RoaringBitmap; /// 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; +/// Delta segments one index name may pile up before they are merged into one: +/// each is a separate search per query. +pub const MAX_DELTAS: usize = 8; /// IVF partitions for a table of `rows`: about one per sqrt(rows), at least one. pub fn partitions_for(rows: usize) -> usize { @@ -50,6 +56,126 @@ pub async fn index_names(dataset: &Dataset, vector_column: &str) -> Result, +} + +/// Coverage of `vector_column` by every index built on it, or `None` when it +/// has no index. An index that does not record its fragments is an error: +/// counting its rows as uncovered would have them indexed a second time. +pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result> { + let Some(field) = dataset.schema().field(vector_column) else { + return Ok(None); + }; + let mut deltas: BTreeMap = BTreeMap::new(); + let mut covered = RoaringBitmap::new(); + for index in dataset.load_indices().await?.iter() { + if index.fields != [field.id] { + continue; + } + let Some(bitmap) = &index.fragment_bitmap else { + anyhow::bail!( + "index {} does not record the fragments it covers", + index.name + ); + }; + *deltas.entry(index.name.clone()).or_default() += 1; + covered |= bitmap; + } + if deltas.is_empty() { + return Ok(None); + } + let (mut indexed_rows, mut unindexed_rows) = (0, 0); + for fragment in dataset.get_fragments() { + let rows = fragment.count_rows(None).await? as u64; + if covered.contains(fragment.id() as u32) { + indexed_rows += rows; + } else { + unindexed_rows += rows; + } + } + Ok(Some(Coverage { + indexed_rows, + unindexed_rows, + deltas, + })) +} + +/// How to extend an index name that already has `deltas` segments: another +/// delta while there is room, and once the limit is reached every delta +/// merged with the new rows into one. Lance's default options promise no +/// merge, so the count is explicit. +pub fn optimize_policy(deltas: usize) -> OptimizeOptions { + if deltas >= MAX_DELTAS { + OptimizeOptions::merge(deltas) + } else { + OptimizeOptions::append() + } +} + +/// What an index was built as, from Lance's statistics. +#[derive(Debug, Clone, PartialEq)] +pub struct IndexSettings { + /// The most partitions any of its deltas has. + pub partitions: usize, + /// Lance's name for the type, such as `IVF_FLAT`. + pub index_type: String, + pub metric: Option, +} + +pub async fn index_settings(dataset: &Dataset, name: &str) -> Result { + let stats: serde_json::Value = serde_json::from_str(&dataset.index_statistics(name).await?) + .with_context(|| format!("parse index statistics for {name}"))?; + let deltas: Vec<&serde_json::Value> = + stats["indices"].as_array().into_iter().flatten().collect(); + Ok(IndexSettings { + partitions: deltas + .iter() + .filter_map(|i| i["num_partitions"].as_u64()) + .max() + .unwrap_or(0) as usize, + index_type: deltas + .first() + .and_then(|i| i["index_type"].as_str()) + .unwrap_or_default() + .to_string(), + metric: deltas + .first() + .and_then(|i| i["metric_type"].as_str()) + .and_then(|m| DistanceType::try_from(m).ok()), + }) +} + +/// Whether an index built with `partitions` should be trained again for a +/// table of `rows`: once the table has grown to twice the partitions it was +/// sized for, the centroids no longer fit it. +pub fn needs_retrain(partitions: usize, rows: usize) -> bool { + partitions_for(rows) >= 2 * partitions.max(1) +} + +/// Whether the index on `vector_column` is due for [`optimize`] with +/// `retrain`: a single IVF_FLAT name the table has outgrown. A sharded build +/// sizes each name for its shard, and other index types keep the settings +/// they were built with. +pub async fn retrain_due(dataset: &Dataset, vector_column: &str) -> Result { + let Some(cover) = coverage(dataset, vector_column).await? else { + return Ok(false); + }; + let [name] = Vec::from_iter(cover.deltas.keys())[..] else { + return Ok(false); + }; + let settings = index_settings(dataset, name).await?; + Ok(settings.index_type == "IVF_FLAT" + && needs_retrain(settings.partitions, dataset.count_rows(None).await?)) +} + /// 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. @@ -60,7 +186,19 @@ pub async fn build( partitions: usize, replace: bool, ) -> Result<()> { - let params = VectorIndexParams::ivf_flat(partitions, distance_metric(dataset, vector_column)); + let metric = distance_metric(dataset, vector_column); + build_with(dataset, vector_column, name, partitions, replace, metric).await +} + +async fn build_with( + dataset: &mut Dataset, + vector_column: &str, + name: &str, + partitions: usize, + replace: bool, + metric: DistanceType, +) -> Result<()> { + let params = VectorIndexParams::ivf_flat(partitions, metric); dataset .create_index( &[vector_column], @@ -75,16 +213,49 @@ pub async fn build( 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. +/// Cover the rows written since the indices on `vector_column` were built. +/// One name takes them, the one with the fewest deltas, under +/// [`optimize_policy`]; the other names keep the fragments they have, so +/// shards stay disjoint. With `retrain` the index is built again over the +/// whole table with partitions for the current row count, in the metric it +/// has; only a single IVF_FLAT name is rebuilt this way. Lance's own retrain +/// does not size the partitions, so the rebuild goes through [`build`]. +/// 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() + let Some(cover) = coverage(dataset, vector_column).await? else { + return Ok(()); }; - options.index_names = Some(index_names(dataset, vector_column).await?); + if retrain { + if cover.deltas.len() > 1 { + anyhow::bail!( + "the vector index is sharded over {} names; each would be rebuilt over the whole table", + cover.deltas.len() + ); + } + let partitions = partitions_for(dataset.count_rows(None).await?); + for name in cover.deltas.keys() { + let settings = index_settings(dataset, name).await?; + if settings.index_type != "IVF_FLAT" { + anyhow::bail!( + "index {name} is {}; only an IVF_FLAT index is retrained here", + settings.index_type + ); + } + let metric = settings + .metric + .unwrap_or_else(|| distance_metric(dataset, vector_column)); + build_with(dataset, vector_column, name, partitions, true, metric).await?; + } + return Ok(()); + } + let (name, deltas) = cover + .deltas + .iter() + .min_by_key(|(name, deltas)| (**deltas, (*name).clone())) + .map(|(name, deltas)| (name.clone(), *deltas)) + .context("no index to extend")?; + let mut options = optimize_policy(deltas); + options.index_names = Some(vec![name]); dataset.optimize_indices(&options).await?; dataset.checkout_latest().await?; Ok(()) @@ -94,6 +265,23 @@ pub async fn optimize(dataset: &mut Dataset, vector_column: &str, retrain: bool) mod tests { use super::*; + #[test] + fn retrain_once_the_table_quadruples() { + assert!(!needs_retrain(224, 50_000)); + assert!(!needs_retrain(224, 150_000)); + assert!(needs_retrain(224, 200_704)); + assert!(needs_retrain(0, 4)); + assert!(!needs_retrain(4096, 100_000_000)); + } + + #[test] + fn deltas_append_until_the_limit_then_merge() { + assert_eq!(optimize_policy(1).num_indices_to_merge, Some(0)); + assert_eq!(optimize_policy(7).num_indices_to_merge, Some(0)); + assert_eq!(optimize_policy(8).num_indices_to_merge, Some(8)); + assert_eq!(optimize_policy(11).num_indices_to_merge, Some(11)); + } + #[test] fn partitions_follow_the_square_root() { assert_eq!(partitions_for(0), 1); diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index bc23076a11..89248c8ac8 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -569,6 +569,51 @@ def test_vector_index_job(vectors, cluster): again = call("/run", "POST", {"timeout_seconds": 60}) assert again["detected_count"] == 0, again + def coverage_now(): + return json.loads(subprocess.run( + [str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test", + "coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout)[0] + + def optimize(expect_detected): + request = urllib.request.Request( + f"{admin}/api/plugin/job-types/lance_optimize_indices/run", method="POST", + data=json.dumps({"timeout_seconds": 110}).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.load(response) + assert result["detected_count"] == expect_detected and result["error_count"] == 0, result + + # A tenth of the index written after the build brings the optimize job + # in, even though it is far below its absolute budget. + optimize(0) + more = {f"m{i}": point() for i in range(200)} + vectors.put_vectors(vectorBucketName=bucket, indexName=index, + vectors=[{"key": key, "data": {"float32": v}} for key, v in more.items()]) + before = coverage_now() + assert before["unindexed_rows"] == 201, before + optimize(1) + after = coverage_now() + assert after["unindexed_rows"] == 0 and after["segments"] == 2, after + + # Five times the rows the index was sized for: the job trains it again + # instead of extending it, and the partitions follow the new row count. + rows.update(more) + for offset in range(0, 4600, 400): + batch = {f"g{offset + i}": point() for i in range(400)} + rows.update(batch) + vectors.put_vectors(vectorBucketName=bucket, indexName=index, + vectors=[{"key": key, "data": {"float32": v}} for key, v in batch.items()]) + optimize(1) + retrained = coverage_now() + assert retrained["unindexed_rows"] == 0 and retrained["segments"] == 1, retrained + assert retrained["partitions"] > after["partitions"], (before, after, retrained) + found = 0 + for _ in range(20): + query = point() + hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index, + queryVector={"float32": query}, topK=10)["vectors"] + found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10)) + assert found >= 180, f"recall after retrain {found}/200" + def test_invalid_credentials(cluster): client = boto3.client("s3vectors", endpoint_url=cluster["endpoint"], aws_access_key_id="unknown",