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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

* weed-vector: --nprobes overrides the probe policy

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

* test/s3vectors: make clean removes weed-vector-cli

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

* seaweed-vector: test the L2 metric under the name CreateIndex stores

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

* seaweed-vector: test that an append-only optimize serves the new row

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test/s3vectors: record PERF_NOISE in the perf report config

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

* test/s3vectors: count a hit tying the k-th best score toward recall

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* admin: lance_vector_index runs in the Lance lane

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

* test/s3vectors: allow float32 rounding at the recall cutoff

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test/s3vectors: count a repeated hit key once toward recall

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test/s3vectors: credit cutoff ties only up to the slots the top-k leaves

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test/s3vectors: assert every filtered hit passes the filter, indexed or not

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>

* weed-vector: test an open after another engine created the table again

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* seaweed-vector: test which vector indices a retrain is due for

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* grafana: panel for vector indexes the Lance worker built

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

* weed-vector: reject --nprobes 0

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Chris LuandClaude Fable 5.1 authored and GitHub committed 2026-11-04 23:30:27 -08:00
1 parent e465880ee8
commit 48df16ea0d
9 files changed
+587 -46

No files matched your search

+1
View File
@@ -6597,6 +6597,7 @@ dependencies = [
"lance-table",
"prometheus",
"reqwest 0.12.28",
"roaring",
"seaweed-common",
"seaweed-worker-core",
"seaweed-worker-sort",
@@ -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<Vec<Inde
.field(vector_column)
.with_context(|| format!("no such column {vector_column}"))?;
let mut by_name: BTreeMap<String, Vec<Option<RoaringBitmap>>> = BTreeMap::new();
let mut partitions: BTreeMap<String, usize> = 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<Vec<Inde
.entry(index.name.clone())
.or_default()
.push(index.fragment_bitmap.clone());
let count = crate::index::segment_partitions(dataset, vector_column, &index.uuid).await?;
let most = partitions.entry(index.name.clone()).or_default();
*most = count.max(*most);
}
let all: RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect();
let mut out = Vec::new();
@@ -58,6 +64,7 @@ pub async fn coverage(dataset: &Dataset, vector_column: &str) -> Result<Vec<Inde
}
}
out.push(IndexCoverage {
partitions: partitions.get(&name).copied().unwrap_or(0),
index_name: name,
segments: bitmaps.len(),
coverage_known,
+189 -11
View File
@@ -31,11 +31,27 @@ pub fn probes_for(partitions: usize) -> 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<usize> {
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<Option<Probes>> {
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<Op
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)));
let count = segment_partitions(dataset, vector_column, &index.uuid).await?;
partitions = Some(count.max(partitions.unwrap_or(0)));
}
Ok(partitions.map(probes_for))
}
@@ -292,6 +302,37 @@ mod tests {
}
}
/// An index built as cosine on a table without metric metadata keeps
/// cosine through a retrain, and gets partitions for the new row count.
#[tokio::test]
async fn retrain_keeps_the_indexs_own_metric() {
use lance::index::DatasetIndexExt;
use lance::index::vector::VectorIndexParams;
use lance_index::IndexType;
use weed_lance_worker::vector_index::index_settings;
let dir = TempDir::new().unwrap();
let mut dataset =
testing::write_dataset(&dir, &testing::vectors(100, DIM), &["a"; 100]).await;
let params = VectorIndexParams::ivf_flat(2, DistanceType::Cosine);
dataset
.create_index(
&["vector"],
IndexType::Vector,
Some("idx".to_string()),
&params,
false,
)
.await
.unwrap();
dataset.checkout_latest().await.unwrap();
testing::append_dataset(&mut dataset, &testing::vectors(800, DIM), &["b"; 800]).await;
optimize(&mut dataset, "vector", true).await.unwrap();
let settings = index_settings(&dataset, "idx").await.unwrap();
assert_eq!(settings.metric, Some(DistanceType::Cosine));
assert_eq!(settings.index_type, "IVF_FLAT");
assert_eq!(settings.partitions, partitions_for(900));
}
#[tokio::test]
async fn optimize_appends_then_retrain_folds_segments() {
let dir = TempDir::new().unwrap();
@@ -335,6 +376,143 @@ mod tests {
}
}
/// Append and optimize nine times: the deltas climb to the limit and then
/// fold into one, and nothing is ever uncovered after an optimize.
#[tokio::test]
async fn deltas_stay_under_the_limit_across_cycles() {
use weed_lance_worker::vector_index::{MAX_DELTAS, coverage as covered};
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 seen = Vec::new();
for round in 0..9 {
testing::append_dataset(&mut dataset, &testing::vectors(20, DIM), &["b"; 20]).await;
optimize(&mut dataset, &table.vector_column, false)
.await
.unwrap();
let cover = covered(&dataset, "vector").await.unwrap().unwrap();
assert_eq!(cover.unindexed_rows, 0, "round {round}");
seen.push(cover.deltas[INDEX_NAME]);
}
assert_eq!(seen, vec![2, 3, 4, 5, 6, 7, 8, 1, 2]);
assert!(seen.iter().all(|&d| d <= MAX_DELTAS));
}
/// Two names covering disjoint fragments, as a sharded build leaves them:
/// coverage is their union, and an optimize extends one name only.
#[tokio::test]
async fn shards_stay_disjoint_through_an_optimize() {
use lance::index::DatasetIndexExt;
use lance::index::vector::VectorIndexParams;
use lance_index::IndexType;
use weed_lance_worker::vector_index::coverage as covered;
let dir = TempDir::new().unwrap();
let mut dataset =
testing::write_dataset(&dir, &testing::vectors(60, DIM), &["a"; 60]).await;
testing::append_dataset(&mut dataset, &testing::vectors(60, DIM), &["b"; 60]).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, &params)
.name(name.to_string())
.fragments(vec![fragment])
.await
.unwrap();
dataset.checkout_latest().await.unwrap();
}
let cover = covered(&dataset, "vector").await.unwrap().unwrap();
assert_eq!((cover.indexed_rows, cover.unindexed_rows), (120, 0));
assert_eq!(cover.deltas.len(), 2);
testing::append_dataset(&mut dataset, &testing::vectors(30, DIM), &["c"; 30]).await;
let table = VectorTable::resolve(&dataset, None, None).unwrap();
optimize(&mut dataset, &table.vector_column, false)
.await
.unwrap();
let cover = covered(&dataset, "vector").await.unwrap().unwrap();
assert_eq!(cover.unindexed_rows, 0);
assert_eq!(
cover.deltas.values().sum::<usize>(),
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, &params, 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, &params)
.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();
+1
View File
@@ -6518,6 +6518,7 @@ dependencies = [
"lance-table",
"prometheus",
"reqwest 0.12.28",
"roaring",
"seaweed-common",
"seaweed-worker-core",
"seaweed-worker-sort",
+1
View File
@@ -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
+124 -25
View File
@@ -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<Option<u64>> {
/// 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<Option<IndexRows>> {
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<String> = indices.iter().map(|index| index.name.clone()).collect();
names.sort();
names.dedup();
@@ -88,15 +110,27 @@ async fn unindexed_rows(table: &OpenTable) -> Result<Option<u64>> {
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<String, ConfigValue> = 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);
}
}
@@ -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<String, ConfigValue>) {
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<String, ConfigValue> = 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(&parameters, "estimated_runtime_seconds", 0),
3600
);
}
#[test]
fn a_build_gets_longer_than_the_default_execution_timeout() {
let descriptor = VectorIndexHandler::new(String::new()).descriptor();
+197 -9
View File
@@ -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<Vec<S
Ok(names)
}
/// What the indices on a vector column cover, taken together. Sharded builds
/// give one name per shard and each name misses the others' fragments, so
/// coverage is the union over names, never a single name's statistics.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Coverage {
pub indexed_rows: u64,
pub unindexed_rows: u64,
/// Delta segments per index name.
pub deltas: BTreeMap<String, usize>,
}
/// 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<Option<Coverage>> {
let Some(field) = dataset.schema().field(vector_column) else {
return Ok(None);
};
let mut deltas: BTreeMap<String, usize> = 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<DistanceType>,
}
pub async fn index_settings(dataset: &Dataset, name: &str) -> Result<IndexSettings> {
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<bool> {
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);
+45
View File
@@ -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",