From 2bfdd0b8a094e873d75646c61555192f7bc99034 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 10 Sep 2026 11:50:01 -0400 Subject: [PATCH 01/16] feat(index): accept a vector index a table is too small to train Creating a vector index on a table below the quantizer's row floor failed, and an already-trained index whose table later shrank below it blocked `optimize_indices` outright, leaving drop-and-recreate as the only way out. Implements the behaviour specified in #4034 for vector indices, and carries over the empty-index contract #3940 set for scalar and FTS: - `build_empty_vector_index` writes a definition with no files instead of returning `not_supported_source`. Only PQ has a row floor: RQ trains from a pair of values and flat storage needs none, so those cover a small table rather than degrade. - `should_train_index` carries the floor, measured in vectors rather than rows. A column of nulls degrades instead of reaching the quantizer, and a multivector row is worth its list length, estimated by sampling, so a hundred rows of ten vectors train. A fragment subset is exempt: it is a caller-driven segment build that supplies its own fragments and often its own IVF model. - `supported_num_partitions` caps a requested partition count at a codebook's worth of vectors each, counted the same way. A partition's centroid is what search ranks to choose partitions, so one built from a handful of vectors prunes nothing while still paying the quantizer's precision loss. The reduction is logged. - `optimize_indices` trains a segment that carries only its definition from the parameters that definition holds, deciding that before the logical index is opened, since opening it needs a file such a segment does not have. One such segment makes the whole column retrain: the logical index is opened by name, so it would be reached whichever segments the caller asks for. - The merge accepts a rebuild that covers nothing only where the caller established beforehand that the column can no longer train and the segment wrote no files to match, so coverage lost for any other reason still fails. Statistics and vector-index detection tolerate a segment with no file. A nearest-neighbour query answers from the table, and `fast_search` finds nothing rather than falling back to a scan. Existing tests that assert the partition count they requested are given the vectors that count needs, so they keep asserting it exactly; tests whose subject is recall keep their fixtures and assert the count their data supports. A deferred index takes its partition count from the data it is materialized against, through the `target_partition_size` that `VectorIndexDetails` already carries. An absolute `num_partitions` is not preserved across the deferral. --- .../python/tests/test_create_empty_index.py | 42 +- rust/lance/src/dataset/tests/dataset_index.rs | 4 +- rust/lance/src/index.rs | 710 +++++++++++++++++- rust/lance/src/index/append.rs | 162 +++- rust/lance/src/index/create.rs | 39 +- rust/lance/src/index/vector.rs | 177 ++++- rust/lance/src/index/vector/ivf/v2.rs | 45 +- rust/lance/src/index/vector/utils.rs | 4 +- 8 files changed, 1089 insertions(+), 94 deletions(-) diff --git a/python/python/tests/test_create_empty_index.py b/python/python/tests/test_create_empty_index.py index 77d4ab034c9..39a2580b2ff 100644 --- a/python/python/tests/test_create_empty_index.py +++ b/python/python/tests/test_create_empty_index.py @@ -31,17 +31,31 @@ def test_create_empty_vector_index(): data = pa.table({"vector": vectors}) dataset = lance.write_dataset(data, "memory://") - # Currently, vector indices with train=False are not supported - try: - dataset.create_index( - "vector", "IVF_PQ", num_partitions=10, num_sub_vectors=8, train=False - ) - # If we get here, the implementation has been added (unexpected for now) - assert False, ( - "Expected NotImplementedError for train=False on vector index, " - "but succeeded" - ) - except NotImplementedError as e: - # Expected error for unimplemented functionality - error_msg = str(e).lower() - assert "not yet implemented" in error_msg or "not implemented" in error_msg + dataset.create_index( + "vector", "IVF_PQ", num_partitions=10, num_sub_vectors=8, train=False + ) + + # Same shape the scalar case reports: listed, covering nothing yet. + indices = dataset.describe_indices() + assert len(indices) == 1 + stats = dataset.stats.index_stats(indices[0].name) + assert stats["num_indexed_rows"] == 0 + assert stats["num_unindexed_rows"] == dataset.count_rows() + + +def test_create_vector_index_below_the_row_floor(): + """A table too small to train the quantizer takes the index anyway.""" + dim = 32 + values = pc.random(100 * dim).cast(pa.float32()) + vectors = pa.FixedSizeListArray.from_arrays(values, dim) + data = pa.table({"vector": vectors}) + dataset = lance.write_dataset(data, "memory://") + + # 100 vectors cannot train a 256-code codebook, and train defaults to True. + dataset.create_index("vector", "IVF_PQ", num_partitions=10, num_sub_vectors=8) + + indices = dataset.describe_indices() + assert len(indices) == 1 + stats = dataset.stats.index_stats(indices[0].name) + assert stats["num_indexed_rows"] == 0 + assert stats["num_unindexed_rows"] == dataset.count_rows() diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index fb59a2f9da8..d0953ffa9d9 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -86,7 +86,7 @@ async fn test_create_index( false, )])); - let float_arr = generate_random_array(512 * dimension as usize); + let float_arr = generate_random_array(2560 * dimension as usize); let vectors = Arc::new( ::try_new_from_values( float_arr, dimension, @@ -667,7 +667,7 @@ async fn test_create_int8_index( false, )])); - let int8_arr = generate_random_int8_array(512 * dimension as usize); + let int8_arr = generate_random_int8_array(2560 * dimension as usize); let vectors = Arc::new( ::try_new_from_values( int8_arr, dimension, diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 415958a438a..ae5adc34174 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -838,7 +838,7 @@ fn validate_segment_index_details(index_name: &str, segments: &[IndexMetadata]) /// /// Older vector segments may not have `VectorIndexDetails` in the manifest, so /// we also recognize them by the legacy monolithic index file name. -fn segment_has_vector_details(segment: &IndexMetadata) -> bool { +pub(crate) fn segment_has_vector_details(segment: &IndexMetadata) -> bool { segment.index_details.as_ref().map_or_else( || { segment @@ -2649,6 +2649,17 @@ async fn collect_regular_indices_statistics( let mut index_uri: Option = None; for meta in metadatas.iter() { + // An index that covers no fragments has no file to load statistics + // from: it carries its definition and nothing else until there is + // enough data to train it. + if meta + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + { + indices_stats.push(serde_json::json!({})); + continue; + } let index_store = Arc::new(LanceIndexStore::from_dataset_for_existing(ds, meta).await?); let index_details = scalar::fetch_index_details(ds, field_path, meta).await?; if index_uri.is_none() { @@ -2985,12 +2996,15 @@ impl DatasetIndexInternalExt for Dataset { .await? .ok_or_else(|| Error::index(format!("Index with id {} does not exist", uuid)))?; - // Check if this is a vector index by looking at the files list - let is_vector_index = if let Some(files) = &index_meta.files { - // If we have file metadata, check if INDEX_FILE_NAME is in the list + // Declared type first, and the legacy file name only for segments that + // predate details: an index awaiting training declares itself a vector + // index and has no file, so a file-based answer would send it to the + // scalar reader. + let is_vector_index = if index_meta.index_details.is_some() { + segment_has_vector_details(&index_meta) + } else if let Some(files) = &index_meta.files { files.iter().any(|f| f.path == INDEX_FILE_NAME) } else { - // Fall back to file existence check for older indices without file metadata let index_dir = self.indice_files_dir(&index_meta)?; let index_file = index_dir .clone() @@ -6034,17 +6048,687 @@ mod tests { let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); let params = VectorIndexParams::ivf_pq(1, 8, 96, DistanceType::L2, 1); - let result = dataset + dataset .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await; + .await + .expect("a table too small to train on still accepts an index"); - assert!(matches!(result, Err(Error::Unprocessable { .. }))); - if let Error::Unprocessable { message, .. } = result.unwrap_err() { - assert_eq!( - message, - "Not enough rows to train PQ. Requires 256 rows but only 100 available", - ) + // The definition is there and covers nothing: 100 rows cannot train a + // 256-code quantizer, so there is nothing to cover yet. + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + // Every fragment is still unindexed, which is what optimize will pick + // up once the column can train. + let unindexed = dataset.unindexed_fragments("vector_idx").await.unwrap(); + assert_eq!(unindexed.len(), dataset.get_fragments().len()); + } + + /// More partitions requested than the data supports trains fewer of them. + /// + /// A partition wants a codebook's worth of vectors, so 300 vectors support + /// one partition however many are asked for. + #[tokio::test] + async fn test_create_index_with_more_partitions_than_rows() { + let test_dir = tempfile::tempdir().unwrap(); + let rows = 300; + let mut dataset = small_vector_dataset(test_dir.path(), rows).await; + + let params = VectorIndexParams::ivf_pq(1000, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("more partitions than rows should reduce partitions, not fail"); + + // Trained, not degraded: 300 rows clear the 256-code PQ floor. + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + !indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + + // 300 / 256 = 1, not the 1000 requested. + assert_eq!(trained_partitions(&dataset).await, vec![1]); + } + + /// Partition counts for a trained logical vector index, one per segment. + async fn trained_partitions(dataset: &Dataset) -> Vec { + dataset + .open_logical_vector_index("vector", "vector_idx") + .await + .unwrap() + .as_ivf() + .unwrap() + .num_partitions_per_segment() + .into_iter() + .map(|(_, partitions)| partitions) + .collect() + } + + /// Between the floor and a codebook per partition, train fewer partitions. + #[tokio::test] + async fn test_create_index_in_the_reduced_partition_band() { + let test_dir = tempfile::tempdir().unwrap(); + // 8 partitions want 8 * 256 = 2048 vectors; 1000 falls inside the band. + let mut dataset = small_vector_dataset(test_dir.path(), 1000).await; + + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 1000 / 256 = 3. + assert_eq!(trained_partitions(&dataset).await, vec![3]); + } + + /// A codebook's worth of vectors per partition trains the count as asked. + #[tokio::test] + async fn test_create_index_above_the_reduced_partition_band() { + let test_dir = tempfile::tempdir().unwrap(); + // 8 * 256 = 2048, cleared by 3000. + let mut dataset = small_vector_dataset(test_dir.path(), 3000).await; + + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + assert_eq!(trained_partitions(&dataset).await, vec![8]); + } + + /// Build a dataset of `rows` random vectors, one fragment. + async fn small_vector_dataset(dir: &std::path::Path, rows: usize) -> Dataset { + let dimensions = 16; + let field = Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dimensions, + ), + false, + ); + let schema = Arc::new(Schema::new(vec![field])); + let values = generate_random_array(rows * dimensions as usize); + let vectors = + arrow_array::FixedSizeListArray::try_new_from_values(values, dimensions).unwrap(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + Dataset::write(reader, dir.to_str().unwrap(), None) + .await + .unwrap() + } + + /// A dataset of `rows` rows where only the first `non_null` hold a vector. + async fn partly_null_vector_dataset( + dir: &std::path::Path, + rows: usize, + non_null: usize, + ) -> Dataset { + let dimensions = 16; + let field = Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dimensions, + ), + true, + ); + let schema = Arc::new(Schema::new(vec![field])); + let mut builder = arrow_array::builder::FixedSizeListBuilder::new( + arrow_array::builder::Float32Builder::new(), + dimensions, + ); + for row in 0..rows { + if row < non_null { + for value in 0..dimensions { + builder.values().append_value(value as f32); + } + builder.append(true); + } else { + for _ in 0..dimensions { + builder.values().append_null(); + } + builder.append(false); + } + } + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(builder.finish())]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + Dataset::write(reader, dir.to_str().unwrap(), None) + .await + .unwrap() + } + + /// The floor counts vectors, not rows. + /// + /// A column that is mostly null has plenty of rows and nothing to train on, + /// so it has to degrade rather than hand too few vectors to the quantizer. + #[tokio::test] + async fn test_create_index_counts_vectors_not_rows() { + let test_dir = tempfile::tempdir().unwrap(); + // 1000 rows clear the floor; the 100 actual vectors do not. + let mut dataset = partly_null_vector_dataset(test_dir.path(), 1000, 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("a mostly-null column must not be handed to the quantizer"); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "100 vectors cannot train a 256-code quantizer" + ); + } + + /// Append `rows` more random vectors as a new fragment. + async fn append_vectors(dir: &std::path::Path, rows: usize) -> Dataset { + let dimensions = 16; + let field = Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dimensions, + ), + false, + ); + let schema = Arc::new(Schema::new(vec![field])); + let values = generate_random_array(rows * dimensions as usize); + let vectors = + arrow_array::FixedSizeListArray::try_new_from_values(values, dimensions).unwrap(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + Dataset::write( + reader, + dir.to_str().unwrap(), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap() + } + + /// The whole point of degrading: the index fills in once the data arrives. + /// + /// A definition-only index has no model to append to, so optimizing has to + /// train it from scratch and then cover every fragment, including the small + /// one that existed before the threshold was met. + #[tokio::test] + async fn test_degraded_index_trains_on_optimize_once_data_arrives() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 8, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Degraded: 100 rows cannot train a 256-code quantizer. + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + + // 100 -> 500 across two fragments, clearing the floor. + let mut dataset = append_vectors(test_dir.path(), 400).await; + assert_eq!(dataset.count_rows(None).await.unwrap(), 500); + assert_eq!(dataset.get_fragments().len(), 2); + + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .expect("a definition-only index must train once the data is there"); + + // Trained and covering: both fragments, nothing left unindexed. + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + let covered = indices[0].fragment_bitmap.as_ref().unwrap(); + assert_eq!( + covered.len() as usize, + dataset.get_fragments().len(), + "every fragment must be covered, including the pre-threshold one" + ); + assert!( + dataset + .unindexed_fragments("vector_idx") + .await + .unwrap() + .is_empty() + ); + } + + /// A vector query works while the index is still only a definition. + /// + /// The index covers no fragments, so every row is unindexed and the search + /// has to answer from the table itself rather than opening an index file + /// that was never written. + #[tokio::test] + async fn test_vector_query_against_a_definition_only_index() { + let test_dir = tempfile::tempdir().unwrap(); + let rows = 100; + let mut dataset = small_vector_dataset(test_dir.path(), rows).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + assert!( + dataset.load_indices().await.unwrap()[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + + let query = vec![0.0_f32; 16]; + let results = dataset + .scan() + .nearest("vector", &Float32Array::from(query), 5) + .unwrap() + .try_into_batch() + .await + .expect("a query must not open an index file that was never written"); + assert_eq!(results.num_rows(), 5); + } + + /// `fast_search` against an index that is still a definition finds nothing. + /// + /// It restricts the search to what the index covers, and a definition + /// covers no rows, so there is nothing to return rather than a fallback + /// scan of the table. + #[tokio::test] + async fn test_fast_search_against_a_definition_only_index() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + let query = Float32Array::from(vec![0.0_f32; 16]); + let results = dataset + .scan() + .nearest("vector", &query, 5) + .unwrap() + .fast_search() + .try_into_batch() + .await + .expect("fast_search must not fail on an index that covers nothing"); + assert_eq!(results.num_rows(), 0); + } + + /// Statistics for an index that is still a definition report no coverage. + /// + /// The same shape the scalar side reports: the index is listed, indexed + /// rows are zero, and every row counts as unindexed. + #[tokio::test] + async fn test_statistics_for_a_definition_only_index() { + let test_dir = tempfile::tempdir().unwrap(); + let rows = 100; + let mut dataset = small_vector_dataset(test_dir.path(), rows).await; + + let params = VectorIndexParams::ivf_pq(10, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + let stats: serde_json::Value = + serde_json::from_str(&dataset.index_statistics("vector_idx").await.unwrap()).unwrap(); + assert_eq!(stats["num_indexed_rows"].as_u64(), Some(0)); + assert_eq!(stats["num_unindexed_rows"].as_u64(), Some(rows as u64)); + } + + /// An index type without a codebook trains on a table a codebook could not. + /// + /// Only PQ needs a vector per code; RQ needs a pair of values and flat + /// storage needs none, so 100 vectors are enough for these to cover the + /// table rather than degrade to a definition. + #[rstest] + #[case::ivf_flat(VectorIndexParams::ivf_flat(1, DistanceType::L2))] + #[case::ivf_rq(VectorIndexParams::ivf_rq(1, 8, DistanceType::L2))] + #[tokio::test] + async fn test_create_index_without_a_codebook_trains_on_a_small_table( + #[case] params: VectorIndexParams, + ) { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("an index with no codebook has no row floor to clear"); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + !indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "should have trained and covered the table" + ); + } + + /// A multivector row holds many vectors, and the floor counts vectors. + /// + /// 100 rows of 10 vectors give the quantizer 1,000 to train from, so this + /// must train rather than defer on the row count. + #[tokio::test] + async fn test_multivector_floor_counts_vectors_not_rows() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = multivector_dataset(test_dir.path(), 100, 10).await; + + // Multivector columns are cosine-only. + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + !indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "1,000 vectors across 100 rows are enough to train" + ); + } + + /// A table of `rows`, each holding `vectors_per_row` vectors. + async fn multivector_dataset( + dir: &std::path::Path, + rows: usize, + vectors_per_row: usize, + ) -> Dataset { + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, ListBuilder}; + + let dimensions = 16; + let field = Field::new( + "vector", + DataType::List(Arc::new(Field::new( + "item", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dimensions, + ), + true, + ))), + true, + ); + let schema = Arc::new(Schema::new(vec![field])); + + let mut builder = + ListBuilder::new(FixedSizeListBuilder::new(Float32Builder::new(), dimensions)); + for row in 0..rows { + for vector in 0..vectors_per_row { + for value in 0..dimensions { + builder + .values() + .values() + .append_value((row + vector + value as usize) as f32 + 0.5); + } + builder.values().append(true); + } + builder.append(true); } + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(builder.finish())]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + Dataset::write(reader, dir.to_str().unwrap(), None) + .await + .unwrap() + } + + /// The partition cap counts a multivector row's whole list too. + /// + /// 1,000 vectors support three partitions; counting the 100 rows instead + /// would allow only one. + #[tokio::test] + async fn test_partition_cap_counts_multivector_lists() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = multivector_dataset(test_dir.path(), 100, 10).await; + + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 1,000 / 256 = 3, capped from the 8 requested. + assert_eq!(trained_partitions(&dataset).await, vec![3]); + } + + /// And a multivector table that really is too small defers, not errors. + #[tokio::test] + async fn test_multivector_below_the_floor_defers() { + let test_dir = tempfile::tempdir().unwrap(); + // 10 rows of 2 vectors is 20 — far short of a 256-code codebook. + let mut dataset = multivector_dataset(test_dir.path(), 10, 2).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("a small multivector table should take the index, not fail"); + + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + } + + /// The partition cap counts vectors, so blanks cannot inflate it. + /// + /// 3,000 rows holding 300 vectors support one partition, not the eight a + /// row count would appear to allow. + #[tokio::test] + async fn test_partition_cap_ignores_null_rows() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = partly_null_vector_dataset(test_dir.path(), 3000, 300).await; + + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 300 / 256 = 1. + assert_eq!(trained_partitions(&dataset).await, vec![1]); + } + + /// Deleting every row leaves the index defined. + /// + /// The definition is the user's declaration, not a property of the data, so + /// emptying the table must not withdraw it — otherwise reloading a table + /// silently drops the indexes it was created with. + #[tokio::test] + async fn test_vector_index_survives_deleting_all_rows() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 1024).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 8, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + assert_eq!(dataset.load_indices().await.unwrap().len(), 1); + + dataset.delete("true").await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 0); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1, "the index definition must survive"); + + // And it is still there after a reopen, so it lives in the manifest + // rather than in whatever the session happened to hold. + let reopened = Dataset::open(test_dir.path().to_str().unwrap()) + .await + .unwrap(); + assert_eq!(reopened.load_indices().await.unwrap().len(), 1); + } + + /// Retraining an index that is still only a definition does nothing. + /// + /// A definition has no segment to open, and the retrain path resolves that + /// before it reaches for one. + #[tokio::test] + async fn test_retrain_an_index_that_is_still_a_definition() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::retrain()) + .await + .expect("retraining a definition must not look for a segment to open"); + + // Still 100 vectors, so still a definition. + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + } + + /// A table that shrinks below the training threshold keeps working. + /// + /// Optimizing was unrunnable in this state: the quantizer cannot train on + /// what is left, and erroring there blocks index maintenance outright, with + /// dropping and recreating the index as the only way out. + #[tokio::test] + async fn test_optimize_indices_after_shrinking_below_the_threshold() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 1024).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 8, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 1024 -> 224, under a 256-code quantizer's floor. + dataset.delete("_rowid < 800").await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 224); + + // Retrain, not append: append-mode has no new data to index here and so + // never reaches the quantizer, which is the path that fails. + dataset + .optimize_indices(&OptimizeOptions::retrain()) + .await + .expect("optimizing a table that shrank past the threshold must not fail"); + + assert_eq!(dataset.load_indices().await.unwrap().len(), 1); + } + + /// Updating every row leaves the index defined, and optimizing still runs. + #[tokio::test] + async fn test_vector_index_survives_updating_all_rows() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 1024).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 8, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + let update = crate::dataset::UpdateBuilder::new(Arc::new(dataset.clone())) + .set("vector", "vector") + .unwrap() + .build() + .unwrap(); + let updated = update.execute().await.unwrap(); + let mut dataset = updated.new_dataset.as_ref().clone(); + + assert_eq!(dataset.load_indices().await.unwrap().len(), 1); + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .expect("optimizing after a full update must not fail"); + assert_eq!(dataset.load_indices().await.unwrap().len(), 1); + } + + /// A table with no rows at all takes an index, which is the case every + /// other database allows and the one a fresh or reloaded table is in. + #[tokio::test] + async fn test_create_vector_index_on_empty_table() { + let test_dir = tempfile::tempdir().unwrap(); + let dimensions = 16; + let field = Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dimensions, + ), + false, + ); + let schema = Arc::new(Schema::new(vec![field])); + let reader = RecordBatchIterator::new( + Vec::>::new(), + schema.clone(), + ); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let params = VectorIndexParams::ivf_pq(1, 8, 96, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("an empty table still accepts an index"); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + + // Reopening has to find the same definition: it is carried by the + // manifest, not by a file on disk. + let reopened = Dataset::open(test_dir.path().to_str().unwrap()) + .await + .unwrap(); + let indices = reopened.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].name, "vector_idx"); } #[tokio::test] diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 245c861de72..1272fc0fcb3 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -32,7 +32,7 @@ use super::vector::ivf::{ SteadyStateRebalance, VectorSegmentCompatibility, index_type_for_segmented_optimize, optimize_vector_indices, select_steady_state_rebalance, vector_segment_compatibility, }; -use super::vector::{LogicalVectorIndex, fresh_vector_segment_params}; +use super::vector::{LogicalVectorIndex, VectorIndexParams, details, fresh_vector_segment_params}; use super::{CreateIndexBuilder, DatasetIndexInternalExt}; use crate::dataset::Dataset; use crate::dataset::index::LanceIndexStoreExt; @@ -614,6 +614,12 @@ async fn merge_scalar_indices<'a>( } async fn metadata_is_vector_index(dataset: &Dataset, index: &IndexMetadata) -> Result { + // Declared type first: an index awaiting training declares itself a vector + // index and has no file to find it by. + if index.index_details.is_some() { + return Ok(crate::index::segment_has_vector_details(index)); + } + if let Some(files) = &index.files { return Ok(files.iter().any(|file| file.path == INDEX_FILE_NAME)); } @@ -658,33 +664,54 @@ pub async fn merge_indices<'a>( .await } -async fn build_fresh_vector_segment( +/// Whether a rebuild of `params` should expect the definition alone, because +/// the column does not hold the vectors its quantizer needs. +async fn expects_definition_only( + dataset: &Dataset, + field_path: &str, + params: &VectorIndexParams, +) -> Result { + let Some(minimum) = super::vector::vector_quantizer_minimum_rows(¶ms.stages) else { + return Ok(false); + }; + Ok(!super::vector::has_vectors_to_train(dataset, field_path, minimum).await?) +} + +/// Retrain a whole segment over `fragment_bitmap` from index parameters alone. +async fn rebuild_vector_segment( dataset: &Dataset, - logical_index: &LogicalVectorIndex, + name: &str, + params: &VectorIndexParams, field_path: &str, fragment_bitmap: &RoaringBitmap, progress: Arc, ) -> Result { - let (reference_metadata, reference_index) = logical_index.iter().last().ok_or_else(|| { - Error::index(format!( - "Optimize vector index: logical index '{}' has no physical segments", - logical_index.name() - )) - })?; - let params = fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; let mut build_dataset = dataset.clone(); - CreateIndexBuilder::new( - &mut build_dataset, - &[field_path], - IndexType::Vector, - ¶ms, - ) - .name(logical_index.name().to_string()) - .replace(true) - .fragments(fragment_bitmap.iter().collect()) - .progress(progress) - .execute_uncommitted() - .await + CreateIndexBuilder::new(&mut build_dataset, &[field_path], IndexType::Vector, params) + .name(name.to_string()) + .replace(true) + .fragments(fragment_bitmap.iter().collect()) + .progress(progress) + .execute_uncommitted() + .await +} + +/// True when a segment carries only its definition, so there is nothing to open +/// and nothing to append to. +/// +/// Both fields have to say so themselves: an absent file list or bitmap records +/// that the segment was never measured, not that it is empty, and such a +/// segment may still have an index file on disk. +fn is_definition_only_segment(metadata: &IndexMetadata) -> bool { + let wrote_no_files = metadata + .files + .as_ref() + .is_some_and(|files| files.is_empty()); + let covers_nothing = metadata + .fragment_bitmap + .as_ref() + .is_some_and(RoaringBitmap::is_empty); + wrote_no_files && covers_nothing } async fn scan_vector_fragments( @@ -709,12 +736,38 @@ fn fresh_vector_segment_result<'a>( segment: IndexMetadata, expected_fragment_bitmap: &RoaringBitmap, removed_indices: Vec<&'a IndexMetadata>, + definition_expected: bool, ) -> Result> { let fragment_bitmap = segment.fragment_bitmap.ok_or_else(|| { Error::index( "Optimize vector index: newly built segment has no fragment bitmap".to_string(), ) })?; + // A column with too few vectors to train the quantizer yields the + // definition alone: its rows return to unindexed, which is the state a + // table below the threshold belongs in. Accepted only when the caller + // established that beforehand and the segment wrote no files to match, so + // coverage lost for any other reason still fails below. + let wrote_no_files = segment.files.as_ref().is_some_and(|files| files.is_empty()); + if definition_expected && wrote_no_files && fragment_bitmap.is_empty() { + return Ok(IndexMergeResults { + new_uuid: segment.uuid, + removed_indices, + new_fragment_bitmap: fragment_bitmap, + new_dataset_version: segment.dataset_version, + new_index_version: segment.index_version, + new_index_details: segment + .index_details + .map(|details| details.as_ref().clone()) + .ok_or_else(|| { + Error::index( + "Optimize vector index: rebuilt definition has no index details" + .to_string(), + ) + })?, + files: Vec::new(), + }); + } if fragment_bitmap != *expected_fragment_bitmap { return Err(Error::index(format!( "Optimize vector index: newly built segment covers fragments {:?}, expected {:?}", @@ -843,6 +896,52 @@ pub async fn merge_indices_with_unindexed_frags<'a>( return Ok(None); } + // A segment still awaiting training has no file to open and nothing + // to append to, so the whole column is trained from the parameters + // its definition carries, superseding every old segment. One such + // segment is enough to force that: the logical index is opened by + // name, so it would be reached whichever segments the caller asks + // for, and there is no file behind it. + if old_indices + .iter() + .any(|idx| is_definition_only_segment(idx)) + { + if unindexed.is_empty() { + return Ok(None); + } + let params = old_indices + .last() + .and_then(|metadata| metadata.index_details.as_deref()) + .and_then(details::vector_params_from_details) + .ok_or_else(|| { + Error::index(format!( + "Optimize vector index: index '{}' awaits training but carries no parameters", + old_indices[0].name + )) + })?; + let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone(); + let segment = rebuild_vector_segment( + dataset.as_ref(), + &old_indices[0].name, + ¶ms, + &field_path, + &fragment_bitmap, + options.progress.clone(), + ) + .await?; + // Only a result that looks degraded is worth confirming; a + // rebuild that covered rows needs no count to justify it. + let definition_expected = is_definition_only_segment(&segment) + && expects_definition_only(dataset.as_ref(), &field_path, ¶ms).await?; + return fresh_vector_segment_result( + segment, + &fragment_bitmap, + old_indices.to_vec(), + definition_expected, + ) + .map(Some); + } + let full_logical_index = dataset .open_logical_vector_index(&field_path, &old_indices[0].name) .await?; @@ -872,19 +971,34 @@ pub async fn merge_indices_with_unindexed_frags<'a>( )?; if options.retrain || rebuild_dormant { + let (reference_metadata, reference_index) = + logical_index.iter().last().ok_or_else(|| { + Error::index(format!( + "Optimize vector index: logical index '{}' has no physical segments", + logical_index.name() + )) + })?; + let params = + fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone(); - let segment = build_fresh_vector_segment( + let segment = rebuild_vector_segment( dataset.as_ref(), - &logical_index, + logical_index.name(), + ¶ms, &field_path, &fragment_bitmap, options.progress.clone(), ) .await?; + // Only a result that looks degraded is worth confirming; a + // rebuild that covered rows needs no count to justify it. + let definition_expected = is_definition_only_segment(&segment) + && expects_definition_only(dataset.as_ref(), &field_path, ¶ms).await?; return fresh_vector_segment_result( segment, &fragment_bitmap, old_indices.to_vec(), + definition_expected, ) .map(Some); } diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index b3bd6178932..b79d4a7a260 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -248,14 +248,24 @@ impl<'a> CreateIndexBuilder<'a> { .map(|resolved| resolved.canonical_path.as_str()) .unwrap_or(quoted_column.as_str()); + // A fragment list naming the whole table is a whole-table build, not a + // subset: `effective_vector_fragments` normalizes it to `None` so a + // retrain gets the same row floor as a create. let vector_fragments_for_validation = is_builtin_vector_index(self.index_type, self.params) - .then_some(self.fragments.as_deref()) + .then(|| effective_vector_fragments(self.dataset, self.fragments.as_deref())) .flatten(); + let quantizer_minimum_rows = self + .params + .as_any() + .downcast_ref::() + .and_then(|params| super::vector::vector_quantizer_minimum_rows(¶ms.stages)); let train = should_train_index( self.dataset, self.train, - vector_fragments_for_validation, + vector_fragments_for_validation.as_deref(), + quantizer_minimum_rows, + column, ) .await?; @@ -571,11 +581,11 @@ impl<'a> CreateIndexBuilder<'a> { "unable to cast index extension to vector".to_string(), ))?; + // An extension that has not been trained writes nothing: the + // definition is the whole index until there is data for it. if train { ext.create_index(self.dataset, column, &index_id, self.params) .await?; - } else { - todo!("create empty vector index when train=false"); } // Capture file sizes after vector index creation let index_dir = self.dataset.indices_dir().join(index_id.to_string()); @@ -933,10 +943,21 @@ fn is_builtin_vector_index(index_type: IndexType, params: &dyn IndexParams) -> b && params.as_any().is::() } +/// Whether there is enough data to train, as opposed to recording the +/// definition and training later. +/// +/// A quantizer needs one row per code to train at all. Below that the answer is +/// the same as `train=false`: keep the definition, cover no rows, and leave +/// `optimize_indices` to train it once the column fills. +/// +/// An index type with a quantizer is measured in non-null vectors, since a +/// column of nulls trains nothing; every other type is satisfied by any row. async fn should_train_index( dataset: &Dataset, train: bool, vector_fragments: Option<&[u32]>, + minimum_rows: Option, + column: &str, ) -> Result { if !train { return Ok(false); @@ -946,12 +967,20 @@ async fn should_train_index( return Ok(false); } + // A fragment subset is a caller-driven segment build: the caller chose the + // fragments, often supplies the IVF model, and owns the row math. Only + // whole-table builds fall back to a definition-only index. if let Some(fragment_ids) = vector_fragments { dataset.get_fragments_from_ids(fragment_ids)?; return Ok(true); } - Ok(dataset.count_rows(None).await? > 0) + // Only a quantizer counts vectors; every other index type is satisfied by + // any row at all. + let Some(minimum) = minimum_rows.map(|minimum| minimum.max(1)) else { + return Ok(dataset.count_rows(None).await? > 0); + }; + super::vector::has_vectors_to_train(dataset, column, minimum).await } fn vector_params_have_precomputed_ivf(params: &VectorIndexParams) -> bool { diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 061a2f39d85..546dcc86f66 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -20,12 +20,14 @@ pub mod utils; mod fixture_test; use self::{ivf::*, pq::PQIndex}; +use arrow_array::cast::AsArray; use arrow_array::Array; use arrow_schema::{DataType, Schema}; use builder::{IvfIndexBuilder, VectorIndexBuildSummary}; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::stream; +use futures::TryStreamExt; use lance_core::utils::tempfile::TempStdDir; use lance_file::versions::v1::reader::FileReader as V1FileReader; use lance_index::frag_reuse::CompactFragReuseIndex; @@ -505,6 +507,140 @@ impl IndexParams for VectorIndexParams { } } +/// Whether `column` holds at least `minimum` vectors to train a quantizer from. +/// +/// The metadata-only row total settles anything already short of the floor, so +/// the validity count runs only where its answer can change the verdict. +/// +pub(crate) async fn has_vectors_to_train( + dataset: &Dataset, + column: &str, + minimum: usize, +) -> Result { + let total = dataset.count_rows(None).await?; + if total < minimum && !is_multivector(dataset, column)? { + return Ok(false); + } + Ok(count_trainable_vectors(dataset, column, total, minimum).await? >= minimum) +} + +/// Vectors in `column` available to train from. +/// +/// A multivector row holds a list, so its rows bound the vectors from neither +/// side — ten rows can carry a thousand, and an empty list carries none. Those +/// are summed from the lists themselves; everything else counts the rows that +/// hold a vector. Both callers compare the count with a threshold, so `enough` +/// stops the walk as soon as the answer is settled. +pub(crate) async fn count_trainable_vectors( + dataset: &Dataset, + column: &str, + total_rows: usize, + enough: usize, +) -> Result { + if !is_multivector(dataset, column)? { + return count_non_null_vectors(dataset, column, total_rows).await; + } + + let mut scanner = dataset.scan(); + scanner.project(&[column])?; + let mut batches = scanner.try_into_stream().await?; + let mut vectors = 0_usize; + while let Some(batch) = batches.try_next().await? { + let lists = utils::get_column_from_batch(&batch, column)?; + let lists = lists.as_list::(); + for row in 0..lists.len() { + if !lists.is_null(row) { + vectors = vectors.saturating_add(lists.value_length(row) as usize); + } + } + if vectors >= enough { + return Ok(vectors); + } + } + Ok(vectors) +} + +fn is_multivector(dataset: &Dataset, column: &str) -> Result { + let (vector_type, _) = get_vector_type(dataset.schema(), column)?; + Ok(matches!(vector_type, DataType::List(_))) +} + +/// Count the rows of `column` holding a vector. +/// +/// A column that cannot hold nulls has one per row, so `total_rows` is the +/// answer and nothing is read. +async fn count_non_null_vectors( + dataset: &Dataset, + column: &str, + total_rows: usize, +) -> Result { + let nullable = dataset + .schema() + .field(column) + .is_none_or(|field| field.nullable); + if !nullable { + return Ok(total_rows); + } + + dataset + .count_rows(Some(format!("{column} IS NOT NULL"))) + .await +} + +/// The partition count a whole-table build can train. +/// +/// A partition's centroid is what search ranks to choose partitions, so one +/// built from a handful of vectors prunes nothing while still paying the +/// quantizer's precision loss. The vectors available therefore cap the count, +/// at a codebook's worth each; without a quantizer the only limit is KMeans +/// needing a vector per centroid. +async fn supported_num_partitions( + dataset: &Dataset, + column: &str, + stages: &[StageParams], + requested: usize, + mode: &str, +) -> Result { + let total = dataset.count_rows(None).await?; + let enough = vectors_for_partitions(stages, requested); + let vectors = count_trainable_vectors(dataset, column, total, enough).await?; + if vectors >= enough { + return Ok(requested); + } + + let supported = match vector_quantizer_minimum_rows(stages) { + Some(target) => requested.min(recommended_num_partitions(vectors, target)), + None => requested.min(vectors).max(1), + }; + if supported < requested { + log::warn!( + "{mode}: training {supported} of the {requested} requested IVF partitions; \ + column '{column}' holds {vectors} vectors" + ); + } + Ok(supported) +} + +/// Vectors needed before `partitions` can each train a quantizer of their own. +/// +/// The upper edge of the reduced-partition band: below this a build trains +/// fewer partitions than asked for. +fn vectors_for_partitions(stages: &[StageParams], partitions: usize) -> usize { + partitions.saturating_mul(vector_quantizer_minimum_rows(stages).unwrap_or(1)) +} + +/// Rows a vector index's quantizer needs before it can be trained. +/// +/// One row per code, which is `2^num_bits` for PQ — 256 at the default 8 bits, +/// and not 256 at any other setting. `None` for an index type whose stages name +/// no quantizer with a row floor. +pub(crate) fn vector_quantizer_minimum_rows(stages: &[StageParams]) -> Option { + stages.iter().find_map(|stage| match stage { + StageParams::PQ(pq) => 1_usize.checked_shl(pq.num_bits as u32), + _ => None, + }) +} + /// Prepare the shared build inputs used by both direct local builds and /// staged shard builds. /// @@ -577,7 +713,21 @@ async fn prepare_vector_segment_build( centroids.len() ))); } - (Some(num_partitions), _) => num_partitions, + (Some(num_partitions), Some(_)) => num_partitions, + // A partition needs a codebook's worth of vectors to summarise, and its + // centroid is what search ranks to choose partitions: one built from a + // handful of vectors prunes nothing and still costs the quantizer's + // precision. So the vectors available cap how many partitions a + // whole-table build trains, at the same `num_centroids` per partition + // that KMeans warns below. Without a quantizer the only limit is + // KMeans needing a vector per centroid. + // + // A fragment subset keeps the count it was given: the segments of one + // logical index have to agree on it. + (Some(requested), None) if fragment_ids.is_none() => { + supported_num_partitions(dataset, column, stages, requested, mode).await? + } + (Some(num_partitions), None) => num_partitions, (None, Some(centroids)) => centroids.len(), (None, None) => { let num_rows = match fragment_ids { @@ -1563,25 +1713,20 @@ pub(crate) async fn build_vector_index_incremental( } } -/// Build an empty vector index without training on data +/// Create a vector index that carries its definition and no data. +/// +/// The parameters live in the index's `index_details` and the fragment bitmap +/// is empty, so it covers no rows and has no file to open. `optimize_indices` +/// trains it once the column holds enough vectors. #[instrument(level = "debug", skip_all)] pub(crate) async fn build_empty_vector_index( _dataset: &Dataset, - column: &str, - name: &str, + _column: &str, + _name: &str, _uuid: Uuid, _params: &VectorIndexParams, ) -> Result> { - // For now, return a NotImplementedError to indicate this functionality - // is still being developed - Err(Error::not_supported_source( - format!( - "Creating empty vector indices with train=False is not yet implemented. \ - Index '{}' for column '{}' cannot be created without training.", - name, column - ) - .into(), - )) + Ok(Vec::new()) } #[instrument(level = "debug", skip_all, fields(old_uuid = old_uuid.to_string(), new_uuid = new_uuid.to_string()))] @@ -2341,11 +2486,11 @@ mod tests { let source_uri = format!("{}/source", test_dir.as_str()); let target_uri = format!("{}/target", test_dir.as_str()); - // Create source dataset with vector column (need at least 256 rows for PQ training) + // Ten partitions, each needing a codebook's worth of vectors. let source_reader = lance_datagen::gen_batch() .col("id", array::step::()) .col("vector", array::rand_vec::(32.into())) - .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + .into_reader_rows(RowCount::from(2560), BatchCount::from(1)); let mut source_dataset = Dataset::write(source_reader, &source_uri, None) .await .unwrap(); diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index d4f3134c2d4..dc10195b3fd 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -3028,12 +3028,20 @@ mod tests { // while 20 neighbors provide a useful recall oracle. const PQ_MATRIX_NUM_ROWS: usize = 320; const PQ_MATRIX_K: usize = 20; - // An 8-bit PQ codebook has 256 centroids, so this is the smallest valid - // training fixture shared by the 8-bit and 4-bit runtime cases. - const LIGHTWEIGHT_PQ_ROWS: usize = 256; + const PQ_MATRIX_NUM_BITS: usize = 8; const LIGHTWEIGHT_PQ_PARTITIONS: usize = 2; + // An 8-bit PQ codebook has 256 codes and needs a training vector for each, + // so this is the smallest valid fixture, shared by the 8-bit and 4-bit + // runtime cases. + const LIGHTWEIGHT_PQ_ROWS: usize = 256; const LIGHTWEIGHT_PQ_SUB_VECTORS: usize = 4; + /// Partitions a fixture of `rows` supports at `num_bits`, capped by the + /// count the build asked for. + fn supported_partitions(rows: usize, num_bits: usize, requested: usize) -> usize { + requested.min(rows / (1usize << num_bits)).max(1) + } + lance_testing::define_stage_event_progress!(RecordingProgress, IndexBuildProgress, Result<()>); #[test] @@ -3872,10 +3880,9 @@ mod tests { let expected_index_type = if use_hnsw { "IVF_HNSW_PQ" } else { "IVF_PQ" }; let expected_sub_index = if use_hnsw { "HNSW" } else { "PQ" }; assert_eq!(stats["index_type"], expected_index_type); - assert_eq!( - stats["indices"][0]["num_partitions"], - LIGHTWEIGHT_PQ_PARTITIONS - ); + let expected_partitions = + supported_partitions(LIGHTWEIGHT_PQ_ROWS, num_bits, LIGHTWEIGHT_PQ_PARTITIONS); + assert_eq!(stats["indices"][0]["num_partitions"], expected_partitions); assert_eq!( stats["indices"][0]["sub_index"]["index_type"], expected_sub_index @@ -5603,7 +5610,7 @@ mod tests { } } - fn pq_matrix_batch() -> RecordBatch + fn pq_matrix_batch(rows: usize) -> RecordBatch where T: ArrowPrimitiveType + 'static, T::Native: Copy + 'static, @@ -5614,7 +5621,7 @@ mod tests { .with_seed(Seed(42)) .col("id", array::step::()) .col("vector", array::rand_vec::(Dimension::from(DIM as u32))) - .into_batch_rows(RowCount::from(PQ_MATRIX_NUM_ROWS as u64)) + .into_batch_rows(RowCount::from(rows as u64)) .unwrap() } @@ -5628,7 +5635,7 @@ mod tests { ivf_params.sample_rate = PQ_MATRIX_NUM_ROWS; let pq_params = PQBuildParams { num_sub_vectors: 4, - num_bits: 8, + num_bits: PQ_MATRIX_NUM_BITS, max_iters: 2, sample_rate: 1, ..Default::default() @@ -5648,12 +5655,14 @@ mod tests { let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); - let batch = pq_matrix_batch::(); + let params = pq_matrix_params(nlist, distance_type, version.clone()); + let batch = pq_matrix_batch::(PQ_MATRIX_NUM_ROWS); + let expected_partitions = + supported_partitions(PQ_MATRIX_NUM_ROWS, PQ_MATRIX_NUM_BITS, nlist); let schema = batch.schema(); let query = batch["vector"].as_fixed_size_list().value(0); let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); - let params = pq_matrix_params(nlist, distance_type, version.clone()); dataset .create_index( &["vector"], @@ -5673,7 +5682,7 @@ mod tests { let index = &indices[0]; assert_eq!(index["index_type"], "IVF_PQ"); assert_eq!(index["metric_type"], distance_type.to_string()); - assert_eq!(index["num_partitions"], nlist); + assert_eq!(index["num_partitions"], expected_partitions); assert_eq!(index["sub_index"]["index_type"], "PQ"); assert_eq!( index["index_file_version"], @@ -6051,7 +6060,7 @@ mod tests { #[tokio::test] async fn test_ivf_pq_f64_smoke(#[case] version: IndexFileVersion) { let test_dir = TempStrDir::default(); - let batch = pq_matrix_batch::(); + let batch = pq_matrix_batch::(PQ_MATRIX_NUM_ROWS); let schema = batch.schema(); let vectors = Arc::new(batch["vector"].as_fixed_size_list().clone()); let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); @@ -6872,11 +6881,11 @@ mod tests { #[tokio::test] async fn test_index_stats( #[values( - (VectorIndexParams::ivf_flat(4, DistanceType::Hamming), IndexType::IvfFlat), - (VectorIndexParams::ivf_pq(4, 8, 8, DistanceType::L2, 10), IndexType::IvfPq), + (VectorIndexParams::ivf_flat(2, DistanceType::Hamming), IndexType::IvfFlat), + (VectorIndexParams::ivf_pq(2, 8, 8, DistanceType::L2, 10), IndexType::IvfPq), (VectorIndexParams::with_ivf_hnsw_sq_params( DistanceType::Cosine, - IvfBuildParams::new(4), + IvfBuildParams::new(2), Default::default(), Default::default() ), IndexType::IvfHnswSq), @@ -6887,7 +6896,7 @@ mod tests { let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); - let nlist = 4; + let nlist = 2; let (mut dataset, _) = match params.metric_type { DistanceType::Hamming => generate_test_dataset::(test_uri, 0..2).await, _ => generate_test_dataset::(test_uri, 0.0..1.0).await, diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index 012d8d3a703..94d93cb74ba 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -30,7 +30,7 @@ use crate::{Error, Result}; /// - Simple column names: "column" /// - Nested paths: "parent.child" or "parent.child.grandchild" /// - Backtick-escaped field names: "parent.`field.with.dots`" -fn get_column_from_batch(batch: &RecordBatch, column: &str) -> Result { +pub(crate) fn get_column_from_batch(batch: &RecordBatch, column: &str) -> Result { // Try to get the column directly first (fast path for simple columns) if let Some(col) = batch.column_by_name(column) { return Ok(col.clone()); @@ -85,7 +85,7 @@ fn get_column_from_batch(batch: &RecordBatch, column: &str) -> Result Ok(current_array) } -async fn estimate_multivector_vectors_per_row( +pub(crate) async fn estimate_multivector_vectors_per_row( dataset: &Dataset, column: &str, num_rows: usize, From 23811ab17ba1347e4d463d6afe14621ddffb03b0 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 10 Sep 2026 22:51:03 -0400 Subject: [PATCH 02/16] test(index): size fixtures for the partitions they request Reducing a partition count the data cannot train left several tests asserting numbers their fixtures no longer produce. Each is given the vectors its request needs, so the assertions stand as written rather than being lowered. - `test_ann_with_deletion` asserts a cached-entry count that is per partition, so four partitions need four codebooks' worth of vectors. Its two fragments, the delete that clears the first, and the ids expected afterwards all derive from one row count now, since they have to agree. - `test_ivf_centroids_exposed` asserts one centroid per requested partition. - `test_index_cache_size` and `test_index_cache_size_bytes` size their cache behaviour from 128 partitions. - `test_create_ivf_rq_index` expected creating an index on an emptied table to raise `NotImplementedError`. That is what this branch makes work, so it now asserts the index is created and covers nothing. The multivector test builder takes a per-row list length, which lets uneven and empty lists be built as easily as uniform ones: a row of ten vectors among ninety-nine empty ones counts as ten, and a column of empty lists counts as none. Both defer rather than reaching the quantizer. --- python/python/tests/test_ivf_centroids.py | 5 +- python/python/tests/test_vector_index.py | 33 +++++++------ rust/lance/src/dataset/scanner.rs | 20 +++++--- rust/lance/src/index.rs | 57 ++++++++++++++++++++++- rust/lance/src/index/vector.rs | 4 +- 5 files changed, 94 insertions(+), 25 deletions(-) diff --git a/python/python/tests/test_ivf_centroids.py b/python/python/tests/test_ivf_centroids.py index 01c2463ab62..34827c330ce 100644 --- a/python/python/tests/test_ivf_centroids.py +++ b/python/python/tests/test_ivf_centroids.py @@ -8,7 +8,10 @@ def test_ivf_centroids_exposed(tmp_path): """Verify that centroids for an IVF-based index are exposed via both""" - dim, rows, parts = 4, 256, 8 + # A partition is trained only when the data can give it a 256-code + # codebook's worth of vectors, so the fixture covers every partition. + dim, parts = 4, 8 + rows = parts * 256 vecs = pa.array( np.random.randn(rows, dim).tolist(), type=pa.list_(pa.float32(), dim), diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index eb2c036f76b..1dce77ee518 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1238,18 +1238,19 @@ def test_create_ivf_rq_index(): assert stats["indices"][0]["sub_index"]["num_bits"] == 5 assert stats["indices"][0]["sub_index"]["packed"] is True - with pytest.raises( - NotImplementedError, - match="Creating empty vector indices with train=False is not yet implemented", - ): - ds.delete("id>=0") - ds = ds.create_index( - "vector", - index_type="IVF_RQ", - num_partitions=4, - num_bits=1, - replace=True, - ) + # An emptied table still takes the index; it carries its settings and + # covers nothing until there is data to train on. + ds.delete("id>=0") + ds = ds.create_index( + "vector", + index_type="IVF_RQ", + num_partitions=4, + num_bits=1, + replace=True, + ) + stats = ds.stats.index_stats("vector_idx") + assert stats["num_indexed_rows"] == 0 + assert stats["num_unindexed_rows"] == 0 zero_vectors = np.zeros((1000, 128)).astype(np.float32).tolist() tbl = pa.Table.from_pydict( @@ -1863,7 +1864,9 @@ def query_index(ds, ntimes, q=None): }, ) - tbl = create_table(nvec=1024, ndim=16) + # Each of the 128 partitions is trained only when the data can give it a + # 256-code codebook's worth of vectors. + tbl = create_table(nvec=128 * 256, ndim=16) dataset = lance.write_dataset(tbl, tmp_path / "test") dataset.create_index( @@ -1908,7 +1911,9 @@ def query_index(ds, ntimes, q=None): }, ) - tbl = create_table(nvec=1024, ndim=16) + # Each of the 128 partitions is trained only when the data can give it a + # 256-code codebook's worth of vectors. + tbl = create_table(nvec=128 * 256, ndim=16) dataset = lance.write_dataset(tbl, tmp_path / "test") dataset.create_index( diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8733b167263..9bada17bbea 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -12040,10 +12040,16 @@ mod test { data_storage_version: LanceFileVersion, #[values(false, true)] stable_row_ids: bool, ) { + const PARTITIONS: usize = 4; + // A partition is trained only when the data can give it a codebook's + // worth of vectors, so the fixture has to cover every partition it asks + // for; the cached-entry count below counts per partition. + const ROWS: usize = PARTITIONS * 256; + let vec_params = vec![ // TODO: re-enable diskann test when we can tune to get reproducible results. // VectorIndexParams::with_diskann_params(MetricType::L2, DiskANNParams::new(10, 1.5, 10)), - VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 2), + VectorIndexParams::ivf_pq(PARTITIONS, 8, 2, MetricType::L2, 2), ]; for params in vec_params { use lance_arrow::FixedSizeListArrayExt; @@ -12066,14 +12072,14 @@ mod test { // vectors are [1, 1, 1, ...] [2, 2, 2, ...] let vector_values: Float32Array = - (0..32 * 512).map(|v| (v / 32) as f32 + 1.0).collect(); + (0..32 * ROWS).map(|v| (v / 32) as f32 + 1.0).collect(); let vectors = FixedSizeListArray::try_new_from_values(vector_values, 32).unwrap(); let batches = vec![ RecordBatch::try_new( schema.clone(), vec![ - Arc::new(Int32Array::from_iter_values(0..512)), + Arc::new(Int32Array::from_iter_values(0..ROWS as i32)), Arc::new(vectors.clone()), ], ) @@ -12175,7 +12181,7 @@ mod test { RecordBatch::try_new( schema.clone(), vec![ - Arc::new(Int32Array::from_iter_values(512..1024)), + Arc::new(Int32Array::from_iter_values(ROWS as i32..2 * ROWS as i32)), Arc::new(vectors), ], ) @@ -12205,7 +12211,7 @@ mod test { .await .unwrap(); - dataset.delete("i < 512").await.unwrap(); + dataset.delete(&format!("i < {ROWS}")).await.unwrap(); let mut scan = dataset.scan(); scan.nearest("vec", &key, 5).unwrap(); @@ -12224,7 +12230,9 @@ mod test { let batch = &results[0]; // It should not pick up any results from the first fragment - let expected_i = BTreeSet::from_iter(vec![512, 513, 514, 515, 516]); + let first = ROWS as i32; + let expected_i = + BTreeSet::from_iter(vec![first, first + 1, first + 2, first + 3, first + 4]); let column_i = batch.column_by_name("i").unwrap(); let actual_i: BTreeSet = as_primitive_array::(column_i.as_ref()) .values() diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index ae5adc34174..5ff0a59c5b9 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6465,6 +6465,12 @@ mod tests { rows: usize, vectors_per_row: usize, ) -> Dataset { + multivector_dataset_with(dir, &vec![vectors_per_row; rows]).await + } + + /// A multivector table whose row `i` holds `lengths[i]` vectors, so uneven + /// and empty lists can be built as easily as uniform ones. + async fn multivector_dataset_with(dir: &std::path::Path, lengths: &[usize]) -> Dataset { use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, ListBuilder}; let dimensions = 16; @@ -6484,8 +6490,8 @@ mod tests { let mut builder = ListBuilder::new(FixedSizeListBuilder::new(Float32Builder::new(), dimensions)); - for row in 0..rows { - for vector in 0..vectors_per_row { + for (row, &length) in lengths.iter().enumerate() { + for vector in 0..length { for value in 0..dimensions { builder .values() @@ -6503,6 +6509,53 @@ mod tests { .unwrap() } + /// Uneven lists are counted, not extrapolated from one of them. + /// + /// One row of ten vectors and ninety-nine empty ones hold ten vectors, not + /// a thousand, so this is still short of a 256-code codebook. + #[tokio::test] + async fn test_multivector_floor_counts_sparse_lists() { + let test_dir = tempfile::tempdir().unwrap(); + let mut lengths = vec![0_usize; 100]; + lengths[0] = 10; + let mut dataset = multivector_dataset_with(test_dir.path(), &lengths).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("ten vectors should defer, not fail"); + + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + } + + /// A column of empty lists holds no vectors at all. + #[tokio::test] + async fn test_multivector_floor_counts_empty_lists_as_none() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = multivector_dataset_with(test_dir.path(), &[0; 100]).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("no vectors should defer, not fail"); + + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + ); + } + /// The partition cap counts a multivector row's whole list too. /// /// 1,000 vectors support three partitions; counting the 100 rows instead diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 546dcc86f66..eeb925b5f27 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -20,14 +20,14 @@ pub mod utils; mod fixture_test; use self::{ivf::*, pq::PQIndex}; -use arrow_array::cast::AsArray; use arrow_array::Array; +use arrow_array::cast::AsArray; use arrow_schema::{DataType, Schema}; use builder::{IvfIndexBuilder, VectorIndexBuildSummary}; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::stream; use futures::TryStreamExt; +use futures::stream; use lance_core::utils::tempfile::TempStdDir; use lance_file::versions::v1::reader::FileReader as V1FileReader; use lance_index::frag_reuse::CompactFragReuseIndex; From 60867068cdf752b3cb496425586e7b55f6d1dbb4 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 10 Sep 2026 23:35:57 -0400 Subject: [PATCH 03/16] fix(index): size partitions by the IVF sample rate The partition cap asked for a codebook's worth of vectors per partition, which is 256 only at the default 8 bits. IVF training fits each centroid on `sample_rate` vectors, and KMeans warns below `k * 256`, so the number belongs to the IVF stage: with a 4-bit codebook the cap asked for 16 vectors per partition and kept all 8 requested partitions on a column holding 300 vectors. An index that stores vectors uncompressed has no precision to lose against a thin partition, so KMeans needing a vector per centroid stays its only limit. A definition also could not come back at the count it asked for. Its parameters are reconstructed from the stored details, which record `target_partition_size` but not `num_partitions`, so an index created as IVF-8 on a table too small to train stayed IVF-1 however far the table grew. `lance.ivf.num_partitions` joins the IVF runtime hints that the same path already restores; an auto-sized build records none and keeps auto-sizing. --- rust/lance/src/index.rs | 107 ++++++++++++++++++++----- rust/lance/src/index/vector.rs | 44 +++++++--- rust/lance/src/index/vector/details.rs | 40 +++++++++ rust/lance/src/index/vector/ivf/v2.rs | 24 +++--- 4 files changed, 174 insertions(+), 41 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 5ff0a59c5b9..a485618f315 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -3813,6 +3813,7 @@ mod tests { hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, kmeans::{KMeansParams, train_kmeans}, + pq::builder::PQBuildParams, sq::builder::SQBuildParams, }; use lance_io::{ @@ -6146,8 +6147,85 @@ mod tests { assert_eq!(trained_partitions(&dataset).await, vec![8]); } - /// Build a dataset of `rows` random vectors, one fragment. - async fn small_vector_dataset(dir: &std::path::Path, rows: usize) -> Dataset { + /// The cap counts the vectors IVF fits a centroid on, not codebook entries. + #[tokio::test] + async fn test_partition_cap_ignores_the_codebook_size() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 300).await; + + // A 4-bit codebook holds 16 entries, but a partition still wants the + // 256 vectors IVF training samples for one centroid. + let params = VectorIndexParams::ivf_pq(8, 4, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 300 / 256 = 1. + assert_eq!(trained_partitions(&dataset).await, vec![1]); + } + + /// Sampling fewer vectors per centroid makes more partitions supportable. + #[tokio::test] + async fn test_partition_cap_follows_the_configured_sample_rate() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 300).await; + + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams { + num_partitions: Some(8), + sample_rate: 64, + ..Default::default() + }, + PQBuildParams { + num_sub_vectors: 4, + num_bits: 8, + ..Default::default() + }, + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // 300 / 64 = 4. + assert_eq!(trained_partitions(&dataset).await, vec![4]); + } + + /// A definition keeps the partition count it asked for, so the index the + /// table grows into is the one that was requested. + #[tokio::test] + async fn test_deferred_index_trains_the_requested_partitions() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "100 vectors cannot train a 256-code quantizer, so nothing is covered yet" + ); + + // 3100 vectors clear 8 * 256, so the requested count is trainable. + let mut dataset = append_vectors(test_dir.path(), 3000).await; + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + + assert_eq!(trained_partitions(&dataset).await, vec![8]); + } + + /// A reader over one batch of `rows` random 16-dimensional vectors. + fn vector_reader(rows: usize) -> impl arrow_array::RecordBatchReader + Send + 'static { let dimensions = 16; let field = Field::new( "vector", @@ -6162,8 +6240,12 @@ mod tests { let vectors = arrow_array::FixedSizeListArray::try_new_from_values(values, dimensions).unwrap(); let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - Dataset::write(reader, dir.to_str().unwrap(), None) + RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema) + } + + /// Build a dataset of `rows` random vectors, one fragment. + async fn small_vector_dataset(dir: &std::path::Path, rows: usize) -> Dataset { + Dataset::write(vector_reader(rows), dir.to_str().unwrap(), None) .await .unwrap() } @@ -6237,23 +6319,8 @@ mod tests { /// Append `rows` more random vectors as a new fragment. async fn append_vectors(dir: &std::path::Path, rows: usize) -> Dataset { - let dimensions = 16; - let field = Field::new( - "vector", - DataType::FixedSizeList( - Arc::new(Field::new("item", DataType::Float32, true)), - dimensions, - ), - false, - ); - let schema = Arc::new(Schema::new(vec![field])); - let values = generate_random_array(rows * dimensions as usize); - let vectors = - arrow_array::FixedSizeListArray::try_new_from_values(values, dimensions).unwrap(); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); Dataset::write( - reader, + vector_reader(rows), dir.to_str().unwrap(), Some(WriteParams { mode: WriteMode::Append, diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index eeb925b5f27..92bb3b9a314 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -592,8 +592,7 @@ async fn count_non_null_vectors( /// A partition's centroid is what search ranks to choose partitions, so one /// built from a handful of vectors prunes nothing while still paying the /// quantizer's precision loss. The vectors available therefore cap the count, -/// at a codebook's worth each; without a quantizer the only limit is KMeans -/// needing a vector per centroid. +/// at `vectors_per_partition` each. async fn supported_num_partitions( dataset: &Dataset, column: &str, @@ -608,7 +607,7 @@ async fn supported_num_partitions( return Ok(requested); } - let supported = match vector_quantizer_minimum_rows(stages) { + let supported = match vectors_per_partition(stages) { Some(target) => requested.min(recommended_num_partitions(vectors, target)), None => requested.min(vectors).max(1), }; @@ -626,7 +625,33 @@ async fn supported_num_partitions( /// The upper edge of the reduced-partition band: below this a build trains /// fewer partitions than asked for. fn vectors_for_partitions(stages: &[StageParams], partitions: usize) -> usize { - partitions.saturating_mul(vector_quantizer_minimum_rows(stages).unwrap_or(1)) + partitions.saturating_mul(vectors_per_partition(stages).unwrap_or(1)) +} + +/// The vectors one partition's centroid should be fitted on, or `None` when the +/// index stores vectors uncompressed. +/// +/// A codebook's precision loss is paid on every vector, so a centroid fitted on +/// a handful prunes nothing and still costs it. Without a codebook there is no +/// precision to lose, and KMeans needing a vector per centroid is the only +/// limit. +fn vectors_per_partition(stages: &[StageParams]) -> Option { + vector_quantizer_minimum_rows(stages)?; + // IVF training fits each centroid on `sample_rate` vectors, which is also + // where KMeans starts warning (`k * 256` at the default rate). That number + // belongs to the IVF stage: an 8-bit codebook holds the same 256 entries by + // coincidence, and a 4-bit one asking only 16 vectors per partition would + // train centroids that prune nothing. + Some( + stages + .iter() + .find_map(|stage| match stage { + StageParams::Ivf(ivf) => Some(ivf.sample_rate), + _ => None, + }) + .unwrap_or_else(|| IvfBuildParams::default().sample_rate) + .max(1), + ) } /// Rows a vector index's quantizer needs before it can be trained. @@ -714,13 +739,10 @@ async fn prepare_vector_segment_build( ))); } (Some(num_partitions), Some(_)) => num_partitions, - // A partition needs a codebook's worth of vectors to summarise, and its - // centroid is what search ranks to choose partitions: one built from a - // handful of vectors prunes nothing and still costs the quantizer's - // precision. So the vectors available cap how many partitions a - // whole-table build trains, at the same `num_centroids` per partition - // that KMeans warns below. Without a quantizer the only limit is - // KMeans needing a vector per centroid. + // A partition's centroid is what search ranks to choose partitions, so + // one built from a handful of vectors prunes nothing and still costs + // the quantizer's precision. So the vectors available cap how many + // partitions a whole-table build trains, at the IVF sample rate each. // // A fragment subset keeps the count it was given: the segments of one // logical index have to agree on it. diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 836da47adfa..74310ae11ed 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/rust/lance/src/index/vector/details.rs @@ -102,6 +102,16 @@ pub fn vector_index_details(params: &VectorIndexParams) -> prost_types::Any { if let Some(tps) = ivf.target_partition_size { target_partition_size = tps as u64; } + // The requested partition count is what a rebuild has to start + // from: a definition created for a table too small to train + // must come back at the count it asked for once the table + // grows, and re-derive the cap from the data it then has. + if let Some(num_partitions) = ivf.num_partitions { + runtime_hints.insert( + "lance.ivf.num_partitions".to_string(), + num_partitions.to_string(), + ); + } runtime_hints.insert("lance.ivf.max_iters".to_string(), ivf.max_iters.to_string()); runtime_hints.insert( "lance.ivf.sample_rate".to_string(), @@ -191,6 +201,9 @@ pub fn apply_runtime_hints(hints: &HashMap, params: &mut VectorI for stage in &mut params.stages { match stage { StageParams::Ivf(ivf) => { + if let Some(v) = parse(hints, "lance.ivf.num_partitions") { + ivf.num_partitions = Some(v); + } if let Some(v) = parse(hints, "lance.ivf.max_iters") { ivf.max_iters = v; } @@ -1128,6 +1141,7 @@ mod tests { let params = VectorIndexParams::with_ivf_pq_params( DistanceType::L2, IvfBuildParams { + num_partitions: Some(12), max_iters: 100, sample_rate: 512, shuffle_partition_batches: 2048, @@ -1146,6 +1160,13 @@ mod tests { let any = vector_index_details(¶ms); let details = any.to_msg::().unwrap(); + assert_eq!( + details + .runtime_hints + .get("lance.ivf.num_partitions") + .map(|s| s.as_str()), + Some("12") + ); assert_eq!( details .runtime_hints @@ -1221,6 +1242,7 @@ mod tests { let StageParams::Ivf(ivf) = &restored.stages[0] else { panic!() }; + assert_eq!(ivf.num_partitions, Some(12)); assert_eq!(ivf.max_iters, 100); assert_eq!(ivf.sample_rate, 512); assert_eq!(ivf.shuffle_partition_batches, 2048); @@ -1231,6 +1253,22 @@ mod tests { assert_eq!(pq.max_iters, 75); assert_eq!(pq.sample_rate, 128); assert_eq!(pq.kmeans_redos, 3); + + // An auto-sized build records no count, so restoring leaves the + // partition count to be derived from the data again. + let auto_sized = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::default(), + PQBuildParams::default(), + ); + let auto_details = vector_index_details(&auto_sized) + .to_msg::() + .unwrap(); + assert!( + !auto_details + .runtime_hints + .contains_key("lance.ivf.num_partitions") + ); } #[test] @@ -1382,6 +1420,7 @@ mod tests { // Non-default values so the round-trip actually checks preservation // rather than coincidentally matching defaults. let ivf = IvfBuildParams { + num_partitions: Some(12), max_iters: 100, sample_rate: 512, target_partition_size: Some(2048), @@ -1456,6 +1495,7 @@ mod tests { let StageParams::Ivf(ivf) = &restored.stages[0] else { panic!("first stage should be IVF for combo {:?}", combo); }; + assert_eq!(ivf.num_partitions, Some(12)); assert_eq!(ivf.max_iters, 100); assert_eq!(ivf.sample_rate, 512); assert_eq!(ivf.target_partition_size, Some(2048)); diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index dc10195b3fd..151bd6f5675 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -3028,18 +3028,18 @@ mod tests { // while 20 neighbors provide a useful recall oracle. const PQ_MATRIX_NUM_ROWS: usize = 320; const PQ_MATRIX_K: usize = 20; - const PQ_MATRIX_NUM_BITS: usize = 8; const LIGHTWEIGHT_PQ_PARTITIONS: usize = 2; + const LIGHTWEIGHT_IVF_SAMPLE_RATE: usize = 16; // An 8-bit PQ codebook has 256 codes and needs a training vector for each, // so this is the smallest valid fixture, shared by the 8-bit and 4-bit // runtime cases. const LIGHTWEIGHT_PQ_ROWS: usize = 256; const LIGHTWEIGHT_PQ_SUB_VECTORS: usize = 4; - /// Partitions a fixture of `rows` supports at `num_bits`, capped by the - /// count the build asked for. - fn supported_partitions(rows: usize, num_bits: usize, requested: usize) -> usize { - requested.min(rows / (1usize << num_bits)).max(1) + /// Partitions a fixture of `rows` supports when each centroid is fitted on + /// `sample_rate` vectors, capped by the count the build asked for. + fn supported_partitions(rows: usize, sample_rate: usize, requested: usize) -> usize { + requested.min(rows / sample_rate).max(1) } lance_testing::define_stage_event_progress!(RecordingProgress, IndexBuildProgress, Result<()>); @@ -3851,7 +3851,7 @@ mod tests { let mut ivf_params = IvfBuildParams::new(LIGHTWEIGHT_PQ_PARTITIONS); ivf_params.max_iters = 2; - ivf_params.sample_rate = 16; + ivf_params.sample_rate = LIGHTWEIGHT_IVF_SAMPLE_RATE; let pq_params = lightweight_pq_params_with_bits(num_bits); let expected_num_sub_vectors = pq_params.num_sub_vectors; let params = if use_hnsw { @@ -3880,8 +3880,11 @@ mod tests { let expected_index_type = if use_hnsw { "IVF_HNSW_PQ" } else { "IVF_PQ" }; let expected_sub_index = if use_hnsw { "HNSW" } else { "PQ" }; assert_eq!(stats["index_type"], expected_index_type); - let expected_partitions = - supported_partitions(LIGHTWEIGHT_PQ_ROWS, num_bits, LIGHTWEIGHT_PQ_PARTITIONS); + let expected_partitions = supported_partitions( + LIGHTWEIGHT_PQ_ROWS, + LIGHTWEIGHT_IVF_SAMPLE_RATE, + LIGHTWEIGHT_PQ_PARTITIONS, + ); assert_eq!(stats["indices"][0]["num_partitions"], expected_partitions); assert_eq!( stats["indices"][0]["sub_index"]["index_type"], @@ -5635,7 +5638,7 @@ mod tests { ivf_params.sample_rate = PQ_MATRIX_NUM_ROWS; let pq_params = PQBuildParams { num_sub_vectors: 4, - num_bits: PQ_MATRIX_NUM_BITS, + num_bits: 8, max_iters: 2, sample_rate: 1, ..Default::default() @@ -5657,8 +5660,9 @@ mod tests { let test_uri = test_dir.as_str(); let params = pq_matrix_params(nlist, distance_type, version.clone()); let batch = pq_matrix_batch::(PQ_MATRIX_NUM_ROWS); + // `pq_matrix_params` fits each centroid on the whole fixture. let expected_partitions = - supported_partitions(PQ_MATRIX_NUM_ROWS, PQ_MATRIX_NUM_BITS, nlist); + supported_partitions(PQ_MATRIX_NUM_ROWS, PQ_MATRIX_NUM_ROWS, nlist); let schema = batch.schema(); let query = batch["vector"].as_fixed_size_list().value(0); let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); From 9767011e80e023e23c85381002199ab11343fb2d Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 01:15:14 -0400 Subject: [PATCH 04/16] test(index): cover compaction and repeated append on a definition Compaction rewrites fragments and remaps every index whose coverage they intersect. An index that covers nothing has no file to open, so the test pins that compaction leaves it alone rather than reaching for one. The second test pins that a caller optimizing on a schedule keeps getting a no-op while the table stays too small to train, rather than an error on the second pass. --- rust/lance/src/index.rs | 77 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index a485618f315..048fd85e394 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6224,6 +6224,83 @@ mod tests { assert_eq!(trained_partitions(&dataset).await, vec![8]); } + /// Appending to an index that is still a definition, on a table that still + /// cannot train, is a no-op rather than an error, and stays one when + /// repeated: a caller that optimizes on a schedule relies on the commit + /// each call makes. + #[tokio::test] + async fn test_append_to_a_definition_that_still_cannot_train() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Twice, so the second call sees whatever the first one left behind. + for attempt in 0..2 { + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap_or_else(|error| { + panic!("append {attempt} on a definition must not fail: {error}") + }); + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "100 vectors still cannot train a 256-code quantizer" + ); + } + } + + /// Compaction leaves an index that covers nothing alone. It has no file to + /// remap, and an empty fragment bitmap cannot intersect a rewrite group, so + /// the remapper skips it instead of opening a file that is not there. + #[tokio::test] + async fn test_compaction_skips_a_definition_only_index() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 100).await; + + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // A second fragment, so compaction has two to rewrite into one. + let mut dataset = append_vectors(test_dir.path(), 100).await; + assert_eq!(dataset.get_fragments().len(), 2); + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 1000, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.count_rows(None).await.unwrap(), 200); + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "the index should still be a definition covering nothing" + ); + } + /// A reader over one batch of `rows` random 16-dimensional vectors. fn vector_reader(rows: usize) -> impl arrow_array::RecordBatchReader + Send + 'static { let dimensions = 16; From 611b21b86d7cd856ec47cc7270d685f2548afa9c Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 17 Sep 2026 19:38:23 -0400 Subject: [PATCH 05/16] refactor(index): size a deferred index from the data it trains on MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A definition records the settings `VectorIndexDetails` carries — metric, quantizer, `target_partition_size` — and a partition count is not among them. An index that trains later is sized from the data it then holds, so a count asked for on a table too small to train it does not come back. `target_partition_size` is the way to state a shape that outlives the wait, because it is persisted and re-read on every rebuild. --- rust/lance/src/index.rs | 12 ++++---- rust/lance/src/index/vector/details.rs | 40 -------------------------- 2 files changed, 7 insertions(+), 45 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 048fd85e394..a0765cd5d8e 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6193,10 +6193,11 @@ mod tests { assert_eq!(trained_partitions(&dataset).await, vec![4]); } - /// A definition keeps the partition count it asked for, so the index the - /// table grows into is the one that was requested. + /// A definition records no partition count, so the index the table grows + /// into is sized from the data present when it finally trains. A count + /// asked for on a table too small to train it does not survive the wait. #[tokio::test] - async fn test_deferred_index_trains_the_requested_partitions() { + async fn test_deferred_index_sizes_from_the_data_it_trains_on() { let test_dir = tempfile::tempdir().unwrap(); let mut dataset = small_vector_dataset(test_dir.path(), 100).await; @@ -6214,14 +6215,15 @@ mod tests { "100 vectors cannot train a 256-code quantizer, so nothing is covered yet" ); - // 3100 vectors clear 8 * 256, so the requested count is trainable. + // 3100 vectors would clear 8 * 256, but the request was not recorded: + // training derives the count from the data instead, 3100 / 8192 -> 1. let mut dataset = append_vectors(test_dir.path(), 3000).await; dataset .optimize_indices(&OptimizeOptions::default()) .await .unwrap(); - assert_eq!(trained_partitions(&dataset).await, vec![8]); + assert_eq!(trained_partitions(&dataset).await, vec![1]); } /// Appending to an index that is still a definition, on a table that still diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 74310ae11ed..836da47adfa 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/rust/lance/src/index/vector/details.rs @@ -102,16 +102,6 @@ pub fn vector_index_details(params: &VectorIndexParams) -> prost_types::Any { if let Some(tps) = ivf.target_partition_size { target_partition_size = tps as u64; } - // The requested partition count is what a rebuild has to start - // from: a definition created for a table too small to train - // must come back at the count it asked for once the table - // grows, and re-derive the cap from the data it then has. - if let Some(num_partitions) = ivf.num_partitions { - runtime_hints.insert( - "lance.ivf.num_partitions".to_string(), - num_partitions.to_string(), - ); - } runtime_hints.insert("lance.ivf.max_iters".to_string(), ivf.max_iters.to_string()); runtime_hints.insert( "lance.ivf.sample_rate".to_string(), @@ -201,9 +191,6 @@ pub fn apply_runtime_hints(hints: &HashMap, params: &mut VectorI for stage in &mut params.stages { match stage { StageParams::Ivf(ivf) => { - if let Some(v) = parse(hints, "lance.ivf.num_partitions") { - ivf.num_partitions = Some(v); - } if let Some(v) = parse(hints, "lance.ivf.max_iters") { ivf.max_iters = v; } @@ -1141,7 +1128,6 @@ mod tests { let params = VectorIndexParams::with_ivf_pq_params( DistanceType::L2, IvfBuildParams { - num_partitions: Some(12), max_iters: 100, sample_rate: 512, shuffle_partition_batches: 2048, @@ -1160,13 +1146,6 @@ mod tests { let any = vector_index_details(¶ms); let details = any.to_msg::().unwrap(); - assert_eq!( - details - .runtime_hints - .get("lance.ivf.num_partitions") - .map(|s| s.as_str()), - Some("12") - ); assert_eq!( details .runtime_hints @@ -1242,7 +1221,6 @@ mod tests { let StageParams::Ivf(ivf) = &restored.stages[0] else { panic!() }; - assert_eq!(ivf.num_partitions, Some(12)); assert_eq!(ivf.max_iters, 100); assert_eq!(ivf.sample_rate, 512); assert_eq!(ivf.shuffle_partition_batches, 2048); @@ -1253,22 +1231,6 @@ mod tests { assert_eq!(pq.max_iters, 75); assert_eq!(pq.sample_rate, 128); assert_eq!(pq.kmeans_redos, 3); - - // An auto-sized build records no count, so restoring leaves the - // partition count to be derived from the data again. - let auto_sized = VectorIndexParams::with_ivf_pq_params( - DistanceType::L2, - IvfBuildParams::default(), - PQBuildParams::default(), - ); - let auto_details = vector_index_details(&auto_sized) - .to_msg::() - .unwrap(); - assert!( - !auto_details - .runtime_hints - .contains_key("lance.ivf.num_partitions") - ); } #[test] @@ -1420,7 +1382,6 @@ mod tests { // Non-default values so the round-trip actually checks preservation // rather than coincidentally matching defaults. let ivf = IvfBuildParams { - num_partitions: Some(12), max_iters: 100, sample_rate: 512, target_partition_size: Some(2048), @@ -1495,7 +1456,6 @@ mod tests { let StageParams::Ivf(ivf) = &restored.stages[0] else { panic!("first stage should be IVF for combo {:?}", combo); }; - assert_eq!(ivf.num_partitions, Some(12)); assert_eq!(ivf.max_iters, 100); assert_eq!(ivf.sample_rate, 512); assert_eq!(ivf.target_partition_size, Some(2048)); From 6a758e585dc6b43eee314f3fc7a570db9376065c Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 17 Sep 2026 22:35:13 -0400 Subject: [PATCH 06/16] fix(mem_wal): read a maintained index's distance type from its details An index that covers nothing carries its settings with no file behind them, so opening the index to inherit its metric fails and the index cannot be registered as maintained. The metric is a recorded setting, so read it from the index's details and open the index only when those do not decode. A table registers WAL before it holds enough vectors to train, and refusing the definition leaves the index outside the maintained set for the life of the table: the set is a snapshot, so training it later does not add it back. --- rust/lance/src/dataset/mem_wal/api.rs | 86 ++++++++++++++++++++++----- 1 file changed, 72 insertions(+), 14 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index ef55d20d5ab..dde5de98888 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -823,20 +823,31 @@ async fn load_vector_index_config( let column = field.name.clone(); // Inherit the base table's distance type so the in-memory index and the - // base index produce comparable distances. Surface the open error - // instead of silently defaulting to L2 — flushed `IVF_HNSW_SQ` files - // bake this metric into their on-disk metadata, so a wrong default would - // be durable corruption. - let distance_type = dataset - .open_vector_index(&column, &index_meta.uuid, &NoOpMetricsCollector) - .await - .map_err(|e| { - Error::invalid_input(format!( - "Failed to open base vector index '{}' to inherit distance type: {}", - index_name, e - )) - })? - .metric_type(); + // base index produce comparable distances. The index's recorded details + // state it, and for an index that covers nothing they are the only source: + // it carries its settings with no file to open. Opening the index covers + // an entry whose details do not decode. Surface the failure instead of + // silently defaulting to L2 — flushed `IVF_HNSW_SQ` files bake this metric + // into their on-disk metadata, so a wrong default would be durable + // corruption. + let recorded = index_meta + .index_details + .as_deref() + .and_then(crate::index::vector::details::vector_params_from_details) + .map(|params| params.metric_type); + let distance_type = match recorded { + Some(distance_type) => distance_type, + None => dataset + .open_vector_index(&column, &index_meta.uuid, &NoOpMetricsCollector) + .await + .map_err(|e| { + Error::invalid_input(format!( + "Failed to open base vector index '{}' to inherit distance type: {}", + index_name, e + )) + })? + .metric_type(), + }; Ok(match hnsw_params { Some(params) => MemIndexConfig::hnsw_with_params( @@ -1050,6 +1061,53 @@ mod tests { .expect("a Float32 vector column is maintainable"); } + /// An index that covers nothing is maintainable: its recorded details + /// state the distance type, so there is no file to open. + /// + /// A table registers WAL before it holds enough vectors to train, and + /// validation refusing the definition would leave the index outside the + /// maintained set for the life of the table — the set is a snapshot, so + /// training it later does not add it back. + #[tokio::test] + async fn test_validate_maintained_indexes_accepts_a_definition() { + use crate::index::vector::VectorIndexParams; + use lance_linalg::distance::DistanceType; + + let tmp = tempfile::tempdir().unwrap(); + let uri = format!("{}/base", tmp.path().to_str().unwrap()); + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + true, + )])); + let reader = RecordBatchIterator::new(vec![], schema.clone()); + let mut dataset = Dataset::write(reader, &uri, Some(WriteParams::default())) + .await + .unwrap(); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vector_idx".to_string()), + &VectorIndexParams::ivf_pq(4, 8, 2, DistanceType::Cosine, 1), + true, + ) + .await + .unwrap(); + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "an empty table trains nothing, so the index covers no rows" + ); + + validate_maintained_indexes(&dataset, &["vector_idx".to_string()]) + .await + .expect("a definition is maintainable"); + } + #[tokio::test] async fn test_validate_maintained_indexes_accepts_btree() { // Guards the shard-schema plumbing: validation resolves field ids against From 52bcfa35c08ddda796a0560ca8197a7498814ba8 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 17 Sep 2026 23:04:23 -0400 Subject: [PATCH 07/16] fix(index): recognise a definition whose file list reads back as absent The manifest stores an index's file list as a repeated field whose empty case means "sizes unknown", so a segment that recorded an empty list reads back as absent after any commit. Requiring the list to be present stopped a definition being recognised the first time anything committed to the table, and optimizing it then opened an index file that was never written. Read the list as "no files recorded", absent included. Coverage still has to be empty: a segment initialized from another dataset's model holds centroids while covering no fragments, and retraining it would discard them. --- rust/lance/src/index/append.rs | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 1272fc0fcb3..dfc9e2f08ff 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -702,16 +702,32 @@ async fn rebuild_vector_segment( /// Both fields have to say so themselves: an absent file list or bitmap records /// that the segment was never measured, not that it is empty, and such a /// segment may still have an index file on disk. +/// Whether a segment has no index data behind it, so training is the work it +/// is waiting for. +/// +/// Both halves carry weight. Coverage alone is not enough: a segment +/// initialized from another dataset's model holds centroids while covering no +/// fragments, and retraining it from its parameters would discard them. +/// +/// The file list is read as "no files recorded", absent included. The manifest +/// stores it as a repeated field whose empty case means "sizes unknown", so a +/// segment that recorded an empty list reads back as absent after any commit — +/// treating absent as "has files" would stop recognising a definition the first +/// time anything else commits to the table. +/// +/// A segment whose rows were all deleted answers this the same way, and wants +/// the same outcome: nothing can be served from it, and a rebuild is what makes +/// it useful again. fn is_definition_only_segment(metadata: &IndexMetadata) -> bool { - let wrote_no_files = metadata + let no_files_recorded = metadata .files .as_ref() - .is_some_and(|files| files.is_empty()); + .is_none_or(|files| files.is_empty()); let covers_nothing = metadata .fragment_bitmap .as_ref() .is_some_and(RoaringBitmap::is_empty); - wrote_no_files && covers_nothing + no_files_recorded && covers_nothing } async fn scan_vector_fragments( From d6a9711aaee50b8eaccaf4b65eb0874a3846494e Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 17 Sep 2026 23:56:55 -0400 Subject: [PATCH 08/16] docs(index): drop the superseded note on definition-only segments The note said an absent file list means the segment was never measured, which is the reading the same commit stopped relying on. --- rust/lance/src/index/append.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index dfc9e2f08ff..8aa21dcc1a5 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -697,13 +697,7 @@ async fn rebuild_vector_segment( } /// True when a segment carries only its definition, so there is nothing to open -/// and nothing to append to. -/// -/// Both fields have to say so themselves: an absent file list or bitmap records -/// that the segment was never measured, not that it is empty, and such a -/// segment may still have an index file on disk. -/// Whether a segment has no index data behind it, so training is the work it -/// is waiting for. +/// and nothing to append to: training is the work it is waiting for. /// /// Both halves carry weight. Coverage alone is not enough: a segment /// initialized from another dataset's model holds centroids while covering no From 25ffbcd910fc3243f047f967aa392a1e1afd2da4 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 00:09:57 -0400 Subject: [PATCH 09/16] style(index): say what the definition path does, not what it argues Also reflows is_definition_only_segment to satisfy rustfmt. --- rust/lance/src/dataset/mem_wal/api.rs | 8 ++++---- rust/lance/src/index.rs | 7 +++---- rust/lance/src/index/append.rs | 5 +---- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index dde5de98888..1de4d298bb5 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -825,9 +825,9 @@ async fn load_vector_index_config( // Inherit the base table's distance type so the in-memory index and the // base index produce comparable distances. The index's recorded details // state it, and for an index that covers nothing they are the only source: - // it carries its settings with no file to open. Opening the index covers - // an entry whose details do not decode. Surface the failure instead of - // silently defaulting to L2 — flushed `IVF_HNSW_SQ` files bake this metric + // it carries its settings with no file to open. Opening the index is the + // fallback for an entry whose details do not decode. Surface the failure + // rather than silently defaulting to L2 — flushed `IVF_HNSW_SQ` files bake this metric // into their on-disk metadata, so a wrong default would be durable // corruption. let recorded = index_meta @@ -1064,7 +1064,7 @@ mod tests { /// An index that covers nothing is maintainable: its recorded details /// state the distance type, so there is no file to open. /// - /// A table registers WAL before it holds enough vectors to train, and + /// A table can register WAL before it holds enough vectors to train, and /// validation refusing the definition would leave the index outside the /// maintained set for the life of the table — the set is a snapshot, so /// training it later does not add it back. diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index a0765cd5d8e..771a47442e1 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6215,8 +6215,8 @@ mod tests { "100 vectors cannot train a 256-code quantizer, so nothing is covered yet" ); - // 3100 vectors would clear 8 * 256, but the request was not recorded: - // training derives the count from the data instead, 3100 / 8192 -> 1. + // 3100 vectors would clear 8 * 256, but no count was recorded, so + // training derives one from the data: 3100 / 8192 -> 1. let mut dataset = append_vectors(test_dir.path(), 3000).await; dataset .optimize_indices(&OptimizeOptions::default()) @@ -6228,8 +6228,7 @@ mod tests { /// Appending to an index that is still a definition, on a table that still /// cannot train, is a no-op rather than an error, and stays one when - /// repeated: a caller that optimizes on a schedule relies on the commit - /// each call makes. + /// repeated. #[tokio::test] async fn test_append_to_a_definition_that_still_cannot_train() { let test_dir = tempfile::tempdir().unwrap(); diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 8aa21dcc1a5..acf5869a25b 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -713,10 +713,7 @@ async fn rebuild_vector_segment( /// the same outcome: nothing can be served from it, and a rebuild is what makes /// it useful again. fn is_definition_only_segment(metadata: &IndexMetadata) -> bool { - let no_files_recorded = metadata - .files - .as_ref() - .is_none_or(|files| files.is_empty()); + let no_files_recorded = metadata.files.as_ref().is_none_or(|files| files.is_empty()); let covers_nothing = metadata .fragment_bitmap .as_ref() From 7397d20a33256fd4771b74879015debcbecb2433 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 00:18:24 -0400 Subject: [PATCH 10/16] test(index): size the shared vector fixture for the partitions it asks for Two partitions over 400 vectors holds only at a sample rate the fixture covers, and tests built on it assert that topology. --- rust/lance/src/dataset/scanner.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 9bada17bbea..99f0aea31f9 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -7280,7 +7280,7 @@ pub mod test_dataset { use uuid::Uuid; use crate::dataset::WriteParams; - use crate::index::vector::VectorIndexParams; + use crate::index::vector::{StageParams, VectorIndexParams}; // Creates a dataset with 5 batches where each batch has 80 rows // @@ -7379,7 +7379,13 @@ pub mod test_dataset { } pub async fn make_vector_index_with_metric(&mut self, metric: MetricType) -> Result<()> { - let params = VectorIndexParams::ivf_pq(2, 8, 2, metric, 2); + let mut params = VectorIndexParams::ivf_pq(2, 8, 2, metric, 2); + // Two partitions over 400 vectors only holds at a sample rate this + // fixture can cover: the trained count is capped at the vectors + // available per centroid. + if let Some(StageParams::Ivf(ivf)) = params.stages.first_mut() { + ivf.sample_rate = 200; + } self.dataset .create_index( &["vec"], From cebf2e235e76407cf3900dc30bb43f84ff451c2d Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 00:18:24 -0400 Subject: [PATCH 11/16] feat(index): warn when a definition cannot keep a requested partition count The count is not among the settings a definition records, so a caller who asked for one gets a different shape once the index trains. --- rust/lance/src/index/create.rs | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index b79d4a7a260..130ccc7a7c0 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -33,7 +33,7 @@ use lance_index::{ }; use lance_table::format::{IndexMetadata, list_index_files_with_sizes}; use std::{collections::HashMap, future::IntoFuture, sync::Arc}; -use tracing::instrument; +use tracing::{instrument, warn}; use uuid::Uuid; use arrow_array::RecordBatchReader; @@ -269,6 +269,31 @@ impl<'a> CreateIndexBuilder<'a> { ) .await?; + if !train { + // A partition count is not among the settings a definition records, + // so a caller who asked for one gets a different shape when the + // index finally trains. + if let Some(requested) = self + .params + .as_any() + .downcast_ref::() + .and_then(|params| { + params.stages.iter().find_map(|stage| match stage { + StageParams::Ivf(ivf) => ivf.num_partitions, + _ => None, + }) + }) + { + warn!( + column, + requested_num_partitions = requested, + "Recording the index without training it: the partition count \ + will be derived from the data it trains on. Use \ + target_partition_size to state a shape that survives the wait." + ); + } + } + // Load indices from the disk. Names are reserved against every index the // manifest carries: one this build cannot read still owns its name, and // handing that name out again commits two indices under it. From f315b811df99bab6f62e4f480ac03299f8f56e9d Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 10:33:51 -0400 Subject: [PATCH 12/16] feat(index): export the quantizer floor and the trainable-vector count Both answer questions a caller outside the crate has to ask before it decides whether an index can train, and each answered elsewhere is a copy of a rule this crate owns. --- rust/lance/src/index/vector.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 92bb3b9a314..3d29547c52e 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -511,12 +511,7 @@ impl IndexParams for VectorIndexParams { /// /// The metadata-only row total settles anything already short of the floor, so /// the validity count runs only where its answer can change the verdict. -/// -pub(crate) async fn has_vectors_to_train( - dataset: &Dataset, - column: &str, - minimum: usize, -) -> Result { +pub async fn has_vectors_to_train(dataset: &Dataset, column: &str, minimum: usize) -> Result { let total = dataset.count_rows(None).await?; if total < minimum && !is_multivector(dataset, column)? { return Ok(false); @@ -654,14 +649,20 @@ fn vectors_per_partition(stages: &[StageParams]) -> Option { ) } +/// Rows a PQ codebook needs before it can be trained. +/// +/// One row per code, which is `2^num_bits` — 256 at the default 8 bits, and not +/// 256 at any other setting. `None` when the width does not fit a `usize`. +pub fn pq_quantizer_minimum_rows(num_bits: usize) -> Option { + 1_usize.checked_shl(num_bits as u32) +} + /// Rows a vector index's quantizer needs before it can be trained. /// -/// One row per code, which is `2^num_bits` for PQ — 256 at the default 8 bits, -/// and not 256 at any other setting. `None` for an index type whose stages name -/// no quantizer with a row floor. +/// `None` for an index type whose stages name no quantizer with a row floor. pub(crate) fn vector_quantizer_minimum_rows(stages: &[StageParams]) -> Option { stages.iter().find_map(|stage| match stage { - StageParams::PQ(pq) => 1_usize.checked_shl(pq.num_bits as u32), + StageParams::PQ(pq) => pq_quantizer_minimum_rows(pq.num_bits), _ => None, }) } From b5e80be0577114c57e75986c0b30ecbf75025524 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 10:57:33 -0400 Subject: [PATCH 13/16] test(index): fold the repeated floor and partition cases into tables Four multivector floor tests shared a body and three asserted the same outcome, and two partition-band tests differed only in a row count. --- rust/lance/src/index.rs | 166 ++++++++++++---------------------------- 1 file changed, 50 insertions(+), 116 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 771a47442e1..d50ad145f83 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6114,37 +6114,27 @@ mod tests { .collect() } - /// Between the floor and a codebook per partition, train fewer partitions. + /// The cap reduces a requested count only while the data cannot support it. #[tokio::test] - async fn test_create_index_in_the_reduced_partition_band() { - let test_dir = tempfile::tempdir().unwrap(); - // 8 partitions want 8 * 256 = 2048 vectors; 1000 falls inside the band. - let mut dataset = small_vector_dataset(test_dir.path(), 1000).await; - - let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .unwrap(); - - // 1000 / 256 = 3. - assert_eq!(trained_partitions(&dataset).await, vec![3]); - } - - /// A codebook's worth of vectors per partition trains the count as asked. - #[tokio::test] - async fn test_create_index_above_the_reduced_partition_band() { - let test_dir = tempfile::tempdir().unwrap(); - // 8 * 256 = 2048, cleared by 3000. - let mut dataset = small_vector_dataset(test_dir.path(), 3000).await; + async fn test_create_index_partition_cap_follows_the_data() { + // 8 partitions want 8 * 256 = 2048 vectors: 1000 falls inside the band + // where the count is reduced, 3000 clears it. + for (rows, expected) in [(1000, 3), (3000, 8)] { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), rows).await; - let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .unwrap(); + let params = VectorIndexParams::ivf_pq(8, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); - assert_eq!(trained_partitions(&dataset).await, vec![8]); + assert_eq!( + trained_partitions(&dataset).await, + vec![expected], + "{rows} vectors" + ); + } } /// The cap counts the vectors IVF fits a centroid on, not codebook entries. @@ -6577,31 +6567,44 @@ mod tests { ); } - /// A multivector row holds many vectors, and the floor counts vectors. - /// - /// 100 rows of 10 vectors give the quantizer 1,000 to train from, so this - /// must train rather than defer on the row count. + /// A multivector row holds a list, so the floor counts the vectors in the + /// lists rather than the rows that carry them. #[tokio::test] async fn test_multivector_floor_counts_vectors_not_rows() { - let test_dir = tempfile::tempdir().unwrap(); - let mut dataset = multivector_dataset(test_dir.path(), 100, 10).await; + let sparse = { + let mut lengths = vec![0_usize; 100]; + lengths[0] = 10; + lengths + }; + // Lists per row, whether that trains, and the count that decides it. + let cases: [(Vec, bool, &str); 4] = [ + (vec![10; 100], true, "100 rows of 10 vectors give 1,000"), + (vec![2; 10], false, "10 rows of 2 vectors give 20"), + (sparse, false, "one row holding 10 vectors still gives 10"), + (vec![0; 100], false, "empty lists give none"), + ]; - // Multivector columns are cosine-only. - let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .unwrap(); + for (lengths, trains, why) in cases { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = multivector_dataset_with(test_dir.path(), &lengths).await; - let indices = dataset.load_indices().await.unwrap(); - assert_eq!(indices.len(), 1); - assert!( - !indices[0] + // Multivector columns are cosine-only. + let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap_or_else(|error| { + panic!("{why}: a table below the floor takes the index rather than failing: {error}") + }); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1, "{why}"); + let covers_nothing = indices[0] .fragment_bitmap .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), - "1,000 vectors across 100 rows are enough to train" - ); + .is_some_and(roaring::RoaringBitmap::is_empty); + assert_eq!(!covers_nothing, trains, "{why}"); + } } /// A table of `rows`, each holding `vectors_per_row` vectors. @@ -6654,53 +6657,6 @@ mod tests { .unwrap() } - /// Uneven lists are counted, not extrapolated from one of them. - /// - /// One row of ten vectors and ninety-nine empty ones hold ten vectors, not - /// a thousand, so this is still short of a 256-code codebook. - #[tokio::test] - async fn test_multivector_floor_counts_sparse_lists() { - let test_dir = tempfile::tempdir().unwrap(); - let mut lengths = vec![0_usize; 100]; - lengths[0] = 10; - let mut dataset = multivector_dataset_with(test_dir.path(), &lengths).await; - - let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .expect("ten vectors should defer, not fail"); - - let indices = dataset.load_indices().await.unwrap(); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); - } - - /// A column of empty lists holds no vectors at all. - #[tokio::test] - async fn test_multivector_floor_counts_empty_lists_as_none() { - let test_dir = tempfile::tempdir().unwrap(); - let mut dataset = multivector_dataset_with(test_dir.path(), &[0; 100]).await; - - let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .expect("no vectors should defer, not fail"); - - let indices = dataset.load_indices().await.unwrap(); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); - } - /// The partition cap counts a multivector row's whole list too. /// /// 1,000 vectors support three partitions; counting the 100 rows instead @@ -6720,28 +6676,6 @@ mod tests { assert_eq!(trained_partitions(&dataset).await, vec![3]); } - /// And a multivector table that really is too small defers, not errors. - #[tokio::test] - async fn test_multivector_below_the_floor_defers() { - let test_dir = tempfile::tempdir().unwrap(); - // 10 rows of 2 vectors is 20 — far short of a 256-code codebook. - let mut dataset = multivector_dataset(test_dir.path(), 10, 2).await; - - let params = VectorIndexParams::ivf_pq(1, 8, 4, DistanceType::Cosine, 1); - dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) - .await - .expect("a small multivector table should take the index, not fail"); - - let indices = dataset.load_indices().await.unwrap(); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); - } - /// The partition cap counts vectors, so blanks cannot inflate it. /// /// 3,000 rows holding 300 vectors support one partition, not the eight a From 9dd9856b36f4dfcc02eed4293c48f66388816680 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 12:02:01 -0400 Subject: [PATCH 14/16] test(index): cover an append-mode optimize training a definition Append is the mode scheduled maintenance uses, and the only definition test in that mode covered the case where the table still cannot train. --- rust/lance/src/index.rs | 50 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index d50ad145f83..3180ca0aef7 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6114,6 +6114,56 @@ mod tests { .collect() } + /// A table indexed while empty trains on the first append-mode optimize + /// that has data behind it. + /// + /// Writing rows never touches the index, so the optimize is where the + /// definition becomes a real index — and append is the mode scheduled + /// maintenance uses, so it has to be the mode that gets there. + #[tokio::test] + async fn test_append_mode_trains_a_definition_once_data_arrives() { + let test_dir = tempfile::tempdir().unwrap(); + let mut dataset = small_vector_dataset(test_dir.path(), 0).await; + + let params = VectorIndexParams::ivf_pq(2, 8, 4, DistanceType::L2, 1); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + let mut dataset = append_vectors(test_dir.path(), 3000).await; + assert_eq!(dataset.count_rows(None).await.unwrap(), 3000); + let indices = dataset.load_indices().await.unwrap(); + assert!( + indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "writing rows leaves the index a definition" + ); + + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!( + indices.len(), + 1, + "training supersedes the definition rather than adding a delta to it" + ); + assert!( + !indices[0] + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty), + "3,000 vectors clear the floor, so the append trains the column" + ); + // Sized from the 3,000 vectors present, not the 2 the request named. + assert_eq!(trained_partitions(&dataset).await, vec![1]); + } + /// The cap reduces a requested count only while the data cannot support it. #[tokio::test] async fn test_create_index_partition_cap_follows_the_data() { From e6d137812ceda6b459b98501d412f9b7632dd80c Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 12:11:50 -0400 Subject: [PATCH 15/16] refactor(index): name the choice between training and appending Optimize either appends to a trained index or trains one for the first time, and the two share only an entry point. Naming the predicate and lifting the training arm out leaves one early return where there were two. --- rust/lance/src/index/append.rs | 98 ++++++++++++++++++++-------------- 1 file changed, 57 insertions(+), 41 deletions(-) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index acf5869a25b..81dbcda2df1 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -739,6 +739,49 @@ async fn scan_vector_fragments( scanner.try_into_stream().await } +/// Train a column whose index is still only a definition, from the parameters +/// that definition carries. +/// +/// The result supersedes every segment the index had: they hold no data, and +/// what this writes covers the whole column. +async fn train_from_definition<'a>( + dataset: &Arc, + old_indices: &[&'a IndexMetadata], + field_path: &str, + options: &OptimizeOptions, +) -> Result> { + let params = old_indices + .last() + .and_then(|metadata| metadata.index_details.as_deref()) + .and_then(details::vector_params_from_details) + .ok_or_else(|| { + Error::index(format!( + "Optimize vector index: index '{}' awaits training but carries no parameters", + old_indices[0].name + )) + })?; + let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone(); + let segment = rebuild_vector_segment( + dataset.as_ref(), + &old_indices[0].name, + ¶ms, + field_path, + &fragment_bitmap, + options.progress.clone(), + ) + .await?; + // Only a result that looks degraded is worth confirming; a rebuild that + // covered rows needs no count to justify it. + let definition_expected = is_definition_only_segment(&segment) + && expects_definition_only(dataset.as_ref(), field_path, ¶ms).await?; + fresh_vector_segment_result( + segment, + &fragment_bitmap, + old_indices.to_vec(), + definition_expected, + ) +} + fn fresh_vector_segment_result<'a>( segment: IndexMetadata, expected_fragment_bitmap: &RoaringBitmap, @@ -899,54 +942,27 @@ pub async fn merge_indices_with_unindexed_frags<'a>( return Ok(merged); } let rebuild_dormant = live_segments.is_empty() && !dormant_segments.is_empty(); - if rebuild_dormant && unindexed.is_empty() { - return Ok(None); - } - // A segment still awaiting training has no file to open and nothing // to append to, so the whole column is trained from the parameters // its definition carries, superseding every old segment. One such // segment is enough to force that: the logical index is opened by // name, so it would be reached whichever segments the caller asks // for, and there is no file behind it. - if old_indices + let awaits_training = old_indices .iter() - .any(|idx| is_definition_only_segment(idx)) - { - if unindexed.is_empty() { - return Ok(None); - } - let params = old_indices - .last() - .and_then(|metadata| metadata.index_details.as_deref()) - .and_then(details::vector_params_from_details) - .ok_or_else(|| { - Error::index(format!( - "Optimize vector index: index '{}' awaits training but carries no parameters", - old_indices[0].name - )) - })?; - let fragment_bitmap = dataset.fragment_bitmap.as_ref().clone(); - let segment = rebuild_vector_segment( - dataset.as_ref(), - &old_indices[0].name, - ¶ms, - &field_path, - &fragment_bitmap, - options.progress.clone(), - ) - .await?; - // Only a result that looks degraded is worth confirming; a - // rebuild that covered rows needs no count to justify it. - let definition_expected = is_definition_only_segment(&segment) - && expects_definition_only(dataset.as_ref(), &field_path, ¶ms).await?; - return fresh_vector_segment_result( - segment, - &fragment_bitmap, - old_indices.to_vec(), - definition_expected, - ) - .map(Some); + .any(|idx| is_definition_only_segment(idx)); + + // Rebuilding and first training both read the whole column, so + // neither has anything to do until a fragment sits outside the + // index. + if unindexed.is_empty() && (rebuild_dormant || awaits_training) { + return Ok(None); + } + + if awaits_training { + return train_from_definition(&dataset, old_indices, &field_path, options) + .await + .map(Some); } let full_logical_index = dataset From de287c086c7da1e9fd01e4d566c51cb8a3758378 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 18 Sep 2026 12:19:49 -0400 Subject: [PATCH 16/16] test(index): name the covers-nothing check once Fourteen assertions spelled out the same bitmap check inline. --- rust/lance/src/index.rs | 93 +++++++++++------------------------------ 1 file changed, 25 insertions(+), 68 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 3180ca0aef7..0789f4877c1 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -6058,12 +6058,7 @@ mod tests { // 256-code quantizer, so there is nothing to cover yet. let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(segment_covers_nothing(&indices[0])); // Every fragment is still unindexed, which is what optimize will pick // up once the column can train. let unindexed = dataset.unindexed_fragments("vector_idx").await.unwrap(); @@ -6089,18 +6084,22 @@ mod tests { // Trained, not degraded: 300 rows clear the 256-code PQ floor. let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); - assert!( - !indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(!segment_covers_nothing(&indices[0])); // 300 / 256 = 1, not the 1000 requested. assert_eq!(trained_partitions(&dataset).await, vec![1]); } /// Partition counts for a trained logical vector index, one per segment. + /// Whether a segment covers no rows, which is how a definition reads back + /// from the manifest. + fn segment_covers_nothing(index: &IndexMetadata) -> bool { + index + .fragment_bitmap + .as_ref() + .is_some_and(roaring::RoaringBitmap::is_empty) + } + async fn trained_partitions(dataset: &Dataset) -> Vec { dataset .open_logical_vector_index("vector", "vector_idx") @@ -6135,10 +6134,7 @@ mod tests { assert_eq!(dataset.count_rows(None).await.unwrap(), 3000); let indices = dataset.load_indices().await.unwrap(); assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + segment_covers_nothing(&indices[0]), "writing rows leaves the index a definition" ); @@ -6154,10 +6150,7 @@ mod tests { "training supersedes the definition rather than adding a delta to it" ); assert!( - !indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + !segment_covers_nothing(&indices[0]), "3,000 vectors clear the floor, so the append trains the column" ); // Sized from the 3,000 vectors present, not the 2 the request named. @@ -6248,10 +6241,7 @@ mod tests { .unwrap(); let indices = dataset.load_indices().await.unwrap(); assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + segment_covers_nothing(&indices[0]), "100 vectors cannot train a 256-code quantizer, so nothing is covered yet" ); @@ -6291,10 +6281,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + segment_covers_nothing(&indices[0]), "100 vectors still cannot train a 256-code quantizer" ); } @@ -6334,10 +6321,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + segment_covers_nothing(&indices[0]), "the index should still be a definition covering nothing" ); } @@ -6427,10 +6411,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + segment_covers_nothing(&indices[0]), "100 vectors cannot train a 256-code quantizer" ); } @@ -6467,12 +6448,7 @@ mod tests { // Degraded: 100 rows cannot train a 256-code quantizer. let indices = dataset.load_indices().await.unwrap(); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(segment_covers_nothing(&indices[0])); // 100 -> 500 across two fragments, clearing the floor. let mut dataset = append_vectors(test_dir.path(), 400).await; @@ -6518,12 +6494,9 @@ mod tests { .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) .await .unwrap(); - assert!( - dataset.load_indices().await.unwrap()[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(segment_covers_nothing( + &dataset.load_indices().await.unwrap()[0] + )); let query = vec![0.0_f32; 16]; let results = dataset @@ -6609,10 +6582,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); assert!( - !indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty), + !segment_covers_nothing(&indices[0]), "should have trained and covered the table" ); } @@ -6649,10 +6619,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1, "{why}"); - let covers_nothing = indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty); + let covers_nothing = segment_covers_nothing(&indices[0]); assert_eq!(!covers_nothing, trains, "{why}"); } } @@ -6799,12 +6766,7 @@ mod tests { // Still 100 vectors, so still a definition. let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(segment_covers_nothing(&indices[0])); } /// A table that shrinks below the training threshold keeps working. @@ -6896,12 +6858,7 @@ mod tests { let indices = dataset.load_indices().await.unwrap(); assert_eq!(indices.len(), 1); - assert!( - indices[0] - .fragment_bitmap - .as_ref() - .is_some_and(roaring::RoaringBitmap::is_empty) - ); + assert!(segment_covers_nothing(&indices[0])); // Reopening has to find the same definition: it is carried by the // manifest, not by a file on disk.