From 48df16ea0d71a060d1ce980263d1be83318cb04b Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Wed, 4 Nov 2026 23:30:27 -0800 Subject: [PATCH] weed-worker: keep a vector index covering as its table grows (ANN phase 3) (#12414) * seaweed-vector: build and extend the vector index weed_lance_worker::vector_index builds an IVF_FLAT index in the metric the table was created with, sizes partitions from the row count, and extends an index with a new segment or a retrain; the core crate re-exports it and adds the probes a query spends. Tests cover coverage after build, append, retrain and compaction. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector-cli: index build and index optimize Logs go to stderr so stdout stays the JSON a caller parses. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: PERF_INDEX builds the index before the timed queries The cluster fixture hands out its catalog URL, the Makefile builds weed-vector-cli, and the perf suite builds the index after ingest and reports the build time and segment count. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: probe a tenth of the index's partitions, more under a filter A query that names no probe count gets one from the index's partition count. Lance searches at least that many partitions, and under a prefilter keeps going until it has k rows; without one it stops there. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: --nprobes overrides the probe policy Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: test that a build takes the table's metric Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: make clean removes weed-vector-cli Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: measure recall with vectors the table does not hold Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: keep the read cache's doc comment on with_read_cache Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: lance_vector_index builds the index of an S3 Vectors table past a row threshold Detection keeps the tables a vector bucket created, with at least min_rows rows (50,000 by default) and no index on their vector column. Execution builds the IVF_FLAT index in the table's metric with sqrt(rows) partitions. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: optimize only the indices on the vector column Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: resolve the query metric with index::distance_metric Co-Authored-By: Claude Opus 5.5 * seaweed-vector: test the L2 metric under the name CreateIndex stores Co-Authored-By: Claude Opus 5.5 * test/s3vectors: build the index through the admin server and the Lance worker VECTOR_WORKER=1 adds weed admin and weed-worker to the cluster. The new test sets the job's threshold, runs it through the plugin API, checks coverage and recall, and that a row written after the build is found. CI gains the worker leg for both volume servers. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: read an index's partitions, and when a table has outgrown them Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: probes from the shared partition reader, partitions in coverage Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: a vector index reindexes at a tenth of its rows and retrains once the table outgrows it lance_optimize_indices proposes a vector table once its uncovered rows pass min(the absolute budget, a tenth of the indexed rows), and trains the index again instead of extending it when sqrt(rows) has reached twice its partitions. The end-to-end test covers both. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: open a table inside a caller's Lance session Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: keep opened indices between queries in one session Every open used a fresh Lance session, so the index and metadata caches died with the request and each query loaded the IVF index again. --index-cache-mb sizes the shared cache, 512 MiB by default. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: PERF_CLUSTERS draws the perf suite's vectors around shared centres Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: the index test draws embedding-like vectors Independent Gaussian vectors have no locality for an IVF index to use, so recall through one only tracks the fraction of partitions probed; points around shared centres behave like embeddings. The test must also run before the one that stops the engine. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: the optimize and retrain steps draw embedding-like vectors Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: test that compaction drops the deleted key from tied hits Co-Authored-By: Claude Opus 5.5 * seaweed-vector: test that an append-only optimize serves the new row Co-Authored-By: Claude Opus 5.5 * test/s3vectors: record PERF_NOISE in the perf report config Co-Authored-By: Claude Opus 5.5 * seaweed-vector: a filtered query widens to four times the probes, not without bound With no upper bound Lance prepares every partition for the late search even when the minimum probes already hold k rows, and a filtered query through the index cost as much as the exact scan. A query that names its own probe count gets exactly that count. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: report recall through an index without asserting it Co-Authored-By: Claude Opus 5.5 * test/s3vectors: count a hit tying the k-th best score toward recall Co-Authored-By: Claude Opus 5.5 * admin: lance_vector_index runs in the Lance lane Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: check a table's index before counting its rows Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: make build and clean cover the worker leg's binaries Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: a retrain rebuilds the index with partitions for the row count Lance's OptimizeOptions::retrain does not size the partitions: it left a 6,000-row index with one. The retrain goes through the builder with sqrt(rows) partitions, for the job and for weed-vector-cli index optimize --retrain alike. The end-to-end test no longer expects a merge to keep the partition count. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * test/s3vectors: allow float32 rounding at the recall cutoff Co-Authored-By: Claude Opus 5.5 * test/s3vectors: count a repeated hit key once toward recall Co-Authored-By: Claude Opus 5.5 * test/s3vectors: credit cutoff ties only up to the slots the top-k leaves Co-Authored-By: Claude Opus 5.5 * test/s3vectors: assert every filtered hit passes the filter, indexed or not Co-Authored-By: Claude Opus 5.5 * weed-worker: coverage over an index's names, one name extended, a merge at the delta limit A sharded build leaves one name per shard, each missing the others' fragments, so coverage is the union of every index on the column and an optimize gives new rows to the name with the fewest deltas only. At MAX_DELTAS the policy passes OptimizeOptions::merge(deltas); Lance's default options may append instead of merging. A retrain keeps an index's own metric and rebuilds only IVF_FLAT indices. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: the optimize job measures and extends a vector table's indices together Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: tests for bounded deltas, disjoint shards and a retrain's metric Nine append-and-optimize cycles give delta counts 2 to 8 and then 1; two names on disjoint fragments stay disjoint through an optimize; a cosine index on a table without metric metadata is still cosine after a retrain. The end-to-end test expects the appended delta after the first optimize. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: forget the shared session when a table was dropped and created again The session keys index lists and transactions by dataset version, and a table created again at the same location restarts its versions. An open that finds an older version than the newest seen at that URI, or the same version committed later, replaces the session and opens again. Co-Authored-By: Claude Opus 5.5 * weed-vector: test an open after another engine created the table again Co-Authored-By: Claude Opus 5.5 * weed-worker: retrain only a single IVF_FLAT vector index name The job compared each shard name's partitions with the whole table's rows, so a sharded index was always due, and a non-IVF_FLAT index made optimize bail on every run before the new rows were indexed. Co-Authored-By: Claude Opus 5.5 * weed-worker: refuse to retrain a vector index sharded over several names A rebuild per name over the whole table left every shard covering every fragment, so a query searched each row once per shard. Co-Authored-By: Claude Opus 5.5 * seaweed-vector: test which vector indices a retrain is due for Co-Authored-By: Claude Opus 5.5 * weed-vector: replace the session in the same lock that records a recreated table A second open of the recreated table could see the new version recorded and the old session still shared, and keep the dropped table's index list. Co-Authored-By: Claude Opus 5.5 * weed-worker: skip a vector table whose rows cannot be counted Detection returned on the first table whose count_rows failed, so one unreadable table kept every other table from being proposed. It now skips that table like an open or index failure. Co-Authored-By: Claude Opus 5.5 * grafana: panel for vector indexes the Lance worker built Co-Authored-By: Claude Opus 5.5 * seaweed-vector: size default probes from the largest index on the column A sharded build gives the column one index name per shard. The first name alone could be a small shard, and its probe count would then be spent on a larger one. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: skip the probe lookup when a query is given no segments Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: open every table in its own Lance session again One session shared by every open keeps a table's index loaded between queries, but Lance keys what it caches by location and version, so a table dropped and created again at the same location could be served the old table's index list. Telling the two apart from version and commit time needs a map that only grows, a retry loop and clocks that agree across engines. A private session per open, as before, has no such state. Measured at 50k x 384 through an index, eight clients: 164 queries/s with the shared session, 137 without; one client is unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-vector: reject --nprobes 0 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: read the partition count from the opened index Lance's index statistics took 3.6 ms a query on a 200k-row table with 447 partitions. Opening the index and asking it takes 0.4 ms, and the scan that follows finds the index already open. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * ci: give the S3 Vectors worker legs the vendored protoc The worker workspace has no protoc of its own for lance's build scripts, so the worker legs failed at their first build. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: give a vector index build an hour before the admin times it out The job set no execution timeout, so the admin's default of 90 seconds applied, and a build that trains on a large table would be reported as failed and retried while it was still running. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: lance_vector_index proposes at most the admin's max_results Every qualifying table went into one message however many there were. The detector stops at the limit the request carries and sets has_more, as the Go detectors do. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: a vector index job re-checks the table it was proposed for A table emptied since detection completes without a build, and one without a vector column fails instead of reporting it already indexed. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: lance_optimize_indices gets a build's execution timeout A retrain rebuilds the index over the whole table, which the admin's default of 90 seconds would cut off on a large one. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: a vector index run window as long as one build An attempt is capped at the end of the job type's run window plus the drain grace, and the default window is 30 minutes, so a build could be cut off before its own hour was up. The window is an hour too, as the Go iceberg handler sets both. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: a lance_optimize_indices run window as long as one retrain Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: vector index proposals declare a build's runtime A run started by hand while the job type's schedule is off takes no timeout from the descriptor, only the admin's 90 seconds. The admin gives a job at least the runtime its proposal declares, so the build and the optimize proposals of a vector table declare the hour. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: coverage reports no partitions for an index of unknown coverage Reading the partition count goes through Lance's index statistics, which fail on an index that records no fragment bitmap, and the whole coverage report failed with them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * weed-worker: an index that records no fragments is not read as uncovered Its rows counted as uncovered, so detection proposed the table and the execution then failed in Lance's statistics. Coverage fails for such an index instead, and detection skips the table with a warning. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * seaweed-vector: coverage reads partition counts from the opened index Lance's index statistics rewrite the manifest of a table whose fragments carry no row counts, which a coverage report must not do and cannot do with read-only credentials. The report asks each opened segment instead, as the probe lookup does, and the statistics-based reader has no caller left. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Fable 5.1 --- seaweed-vector/Cargo.lock | 1 + seaweed-vector/crates/core/src/coverage.rs | 7 + seaweed-vector/crates/core/src/index.rs | 200 ++++++++++++++++- seaweed-worker/Cargo.lock | 1 + seaweed-worker/crates/lance/Cargo.toml | 1 + .../crates/lance/src/jobs/indices.rs | 149 ++++++++++--- .../crates/lance/src/jobs/vector_index.rs | 23 +- .../crates/lance/src/vector_index.rs | 206 +++++++++++++++++- test/s3vectors/s3vectors_test.py | 45 ++++ 9 files changed, 587 insertions(+), 46 deletions(-) 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",