From a17d7cea537e9e92f652ed10160109736d15a7d8 Mon Sep 17 00:00:00 2001 From: Vivek Date: Sun, 16 Aug 2026 16:11:42 -0700 Subject: [PATCH 1/4] feat(index)!: covering columns for IVF_PQ vector indexes An IVF_PQ index can store the values of chosen extra columns next to its compressed vectors, and a search returns them directly from the index. A query whose projection those columns satisfy no longer reads the base table at all. Each storage format names its own internal columns, so covering detection is a per-storage filter rather than a per-type special case. Nested fields, blob columns, duplicates, reserved storage names and non-IVF_PQ index types are rejected at creation. `fields` is built as `[keyed_id] ++ covering_fields` wherever an index is created, so every covered index satisfies the suffix rule `IndexMetadata::validate_covering_fields` enforces at commit, and segments of one logical index are rejected if they disagree on `covering_fields` -- the read path derives its output schema from the first segment, so a disagreement yields a plan no segment can satisfy. BREAKING CHANGE: `VectorIndexParams` gains a public `covering_columns: Vec` field. The struct is not `#[non_exhaustive]`, so code constructing it with an exhaustive struct literal needs one added line. --- java/lance-jni/src/utils.rs | 1 + rust/lance-index/src/vector/pq/storage.rs | 235 ++- rust/lance-index/src/vector/storage.rs | 130 +- rust/lance-index/src/vector/transform.rs | 173 ++- rust/lance/src/dataset/index.rs | 73 - rust/lance/src/dataset/optimize.rs | 2 + rust/lance/src/dataset/scanner.rs | 25 + rust/lance/src/dataset/tests/dataset_index.rs | 548 ++++--- rust/lance/src/index.rs | 258 ++-- rust/lance/src/index/append.rs | 196 ++- rust/lance/src/index/create.rs | 34 +- rust/lance/src/index/vector.rs | 443 +++++- rust/lance/src/index/vector/builder.rs | 270 +++- rust/lance/src/index/vector/ivf.rs | 31 +- rust/lance/src/index/vector/ivf/v2.rs | 1345 ++++++++++++++++- rust/lance/src/index/vector/utils.rs | 46 + rust/lance/src/io/exec/knn.rs | 773 +++++++++- 17 files changed, 4103 insertions(+), 480 deletions(-) diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index c9d2d2005b1..09bc09207da 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -515,6 +515,7 @@ pub fn get_vector_index_params( version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), }) }, )?; diff --git a/rust/lance-index/src/vector/pq/storage.rs b/rust/lance-index/src/vector/pq/storage.rs index a5ee3496571..0bcd9a3a5ef 100644 --- a/rust/lance-index/src/vector/pq/storage.rs +++ b/rust/lance-index/src/vector/pq/storage.rs @@ -14,7 +14,7 @@ use std::{ use arrow::datatypes::{self, UInt8Type}; use arrow_array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray}; use arrow_array::{ - FixedSizeListArray, RecordBatch, UInt8Array, UInt64Array, + FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, UInt64Array, cast::AsArray, types::{Float32Type, UInt64Type}, }; @@ -43,10 +43,10 @@ use crate::vector::graph::{OrderedFloat, OrderedNode}; use crate::{ INDEX_METADATA_SCHEMA_KEY, IndexMetadata, pb, vector::{ - PQ_CODE_COLUMN, + PART_ID_COLUMN, PQ_CODE_COLUMN, pq::transform::PQTransformer, quantizer::{QuantizerMetadata, QuantizerStorage}, - storage::{DistCalculator, VectorStore}, + storage::{DistCalculator, VectorStore, covering_field_indices_excluding}, transform::Transformer, }, }; @@ -230,16 +230,23 @@ impl ProductQuantizationStorage { transposed: bool, frag_reuse_index: Option>, ) -> Result { - if batch.num_columns() != 2 { - log::warn!( - "PQ storage should have 2 columns, but got {} columns: {}", - batch.num_columns(), - batch.schema(), - ); - batch = batch.project(&[ - batch.schema().index_of(ROW_ID)?, - batch.schema().index_of(PQ_CODE_COLUMN)?, - ])?; + // Require `_rowid` and the PQ code column, but preserve any additional + // ("included"/covering) columns row-aligned. + // Normalize the order to `[_rowid, __pq_code, ]` so the schema + // is stable across builds. `__ivf_part_id` is never a covering column: + // legacy `num_partitions=1` builds (pre-#3606) wrote it into the index + // file, so keep dropping it on load as before. + let row_id_idx = batch.schema().index_of(ROW_ID)?; + let pq_idx = batch.schema().index_of(PQ_CODE_COLUMN)?; + if row_id_idx != 0 || pq_idx != 1 || batch.num_columns() > 2 { + let schema = batch.schema(); + let mut order = Vec::with_capacity(batch.num_columns()); + order.push(row_id_idx); + order.push(pq_idx); + order.extend((0..batch.num_columns()).filter(|&i| { + i != row_id_idx && i != pq_idx && schema.field(i).name() != PART_ID_COLUMN + })); + batch = batch.project(&order)?; } let Some(row_ids) = batch.column_by_name(ROW_ID) else { @@ -285,6 +292,7 @@ impl ProductQuantizationStorage { let transposed_codes = pq_code.values(); let mut new_row_ids = Vec::with_capacity(row_ids.len()); let mut new_codes = Vec::with_capacity(row_ids.len() * num_sub_vectors); + let mut kept_indices: Vec = Vec::with_capacity(row_ids.len()); let row_ids_values = row_ids.values(); for (i, row_id) in row_ids_values.iter().enumerate() { @@ -296,6 +304,7 @@ impl ProductQuantizationStorage { num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } } @@ -311,7 +320,14 @@ impl ProductQuantizationStorage { new_transposed_codes, num_bytes_in_code as i32, )?); - RecordBatch::try_new(batch.schema(), vec![new_row_ids, codes_fsl])? + // Preserve any included/covering columns row-aligned to the kept rows. + rebuild_storage_batch( + batch.schema(), + &batch, + new_row_ids, + codes_fsl, + &UInt32Array::from(kept_indices), + )? }; pq_code = batch[PQ_CODE_COLUMN] .as_fixed_size_list() @@ -529,6 +545,40 @@ where transposed_codes.into() } +/// Rebuild a PQ-storage batch after a row selection (fragment-reuse remap in +/// [`ProductQuantizationStorage::new`] or [`ProductQuantizationStorage::remap`]), +/// preserving any extra ("included"/covering) columns beyond `_rowid` and the PQ +/// code. The already-remapped `row_ids` and transposed `pq_codes` are supplied +/// directly; every other column is gathered from `original` at the surviving +/// positions `kept` so it stays row-aligned. `schema` carries the column order +/// `[_rowid, __pq_code, ]`. +fn rebuild_storage_batch( + schema: SchemaRef, + original: &RecordBatch, + row_ids: ArrayRef, + pq_codes: ArrayRef, + kept: &UInt32Array, +) -> Result { + let columns = schema + .fields() + .iter() + .map(|field| { + let name = field.name().as_str(); + if name == ROW_ID { + Ok(row_ids.clone()) + } else if name == PQ_CODE_COLUMN { + Ok(pq_codes.clone()) + } else { + let col = original.column_by_name(name).ok_or_else(|| { + Error::index(format!("column '{name}' missing from PQ storage batch")) + })?; + Ok(arrow::compute::take(col.as_ref(), kept, None)?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(schema, columns)?) +} + #[async_trait] impl QuantizerStorage for ProductQuantizationStorage { type Metadata = ProductQuantizationMetadata; @@ -610,6 +660,7 @@ impl QuantizerStorage for ProductQuantizationStorage { let transposed_codes = self.pq_code.values(); let mut new_row_ids = Vec::with_capacity(self.len()); let mut new_codes = Vec::with_capacity(self.len() * self.metadata.num_sub_vectors); + let mut kept_indices: Vec = Vec::with_capacity(self.len()); let row_ids = self.row_ids.values(); for (i, row_id) in row_ids.iter().enumerate() { @@ -622,6 +673,7 @@ impl QuantizerStorage for ProductQuantizationStorage { self.metadata.num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } Some(None) => {} None => { @@ -632,6 +684,7 @@ impl QuantizerStorage for ProductQuantizationStorage { self.metadata.num_sub_vectors, i as u32, )); + kept_indices.push(i as u32); } } } @@ -647,7 +700,14 @@ impl QuantizerStorage for ProductQuantizationStorage { new_transposed_codes, num_bytes_in_code as i32, )?); - RecordBatch::try_new(self.schema(), vec![new_row_ids.clone(), codes_fsl])? + // Preserve any included/covering columns row-aligned to the kept rows. + rebuild_storage_batch( + self.schema(), + &self.batch, + new_row_ids.clone(), + codes_fsl, + &UInt32Array::from(kept_indices), + )? }; let transposed_codes = batch[PQ_CODE_COLUMN] .as_fixed_size_list() @@ -722,6 +782,12 @@ impl VectorStore for ProductQuantizationStorage { self.batch.schema_ref() } + /// PQ storage is `[_rowid, __pq_code, ]`, so the covering + /// columns are every field except the row id and the PQ code column. + fn covering_field_indices(&self) -> Vec { + covering_field_indices_excluding(self.schema().as_ref(), &[ROW_ID, PQ_CODE_COLUMN]) + } + fn as_any(&self) -> &dyn std::any::Any { self } @@ -1290,6 +1356,132 @@ mod tests { assert_eq!(storage.row_ids.len(), TOTAL); } + #[tokio::test] + async fn test_build_preserves_extra_column() { + // A covering ("included") column beyond _rowid + __pq_code must be kept + // in the storage batch, row-aligned to the rows, not dropped. + let storage = create_pq_storage_with_extra_column().await; + let batch = storage.batch(); + assert_eq!(batch.num_columns(), 3); + assert!(batch.column_by_name(ROW_ID).is_some()); + assert!(batch.column_by_name(PQ_CODE_COLUMN).is_some()); + let extra = batch + .column_by_name("extra") + .expect("extra column should be preserved") + .as_primitive::(); + // In the fixture extra[i] == row_ids[i] == i, so it must stay aligned. + for (i, row_id) in storage.row_ids().enumerate() { + assert_eq!(extra.value(i) as u64, *row_id); + } + } + + #[tokio::test] + async fn test_build_rejects_mismatched_batch_columns() { + // Batches merged into one storage must carry the same column set. A + // batch with a column the FIRST batch lacks must be rejected: aligning + // to the first batch's schema would silently drop it, losing covering + // data while the index metadata still advertises it. + let codebook = Float32Array::from_iter_values((0..256 * DIM).map(|_| rand::random())); + let codebook = FixedSizeListArray::try_new_from_values(codebook, DIM as i32).unwrap(); + let pq = ProductQuantizer::new(NUM_SUB_VECTORS, 8, DIM, codebook, DistanceType::Dot); + + let vec_field = Field::new( + "vec", + DataType::FixedSizeList( + Field::new_list_field(DataType::Float32, true).into(), + DIM as i32, + ), + true, + ); + let make_fsl = || { + let vectors = Float32Array::from_iter_values((0..TOTAL * DIM).map(|_| rand::random())); + FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap() + }; + let row_ids = || UInt64Array::from_iter_values((0..TOTAL).map(|v| v as u64)); + + let plain_schema = ArrowSchema::new(vec![vec_field.clone(), ROW_ID_FIELD.clone()]); + let plain = RecordBatch::try_new( + plain_schema.into(), + vec![Arc::new(make_fsl()), Arc::new(row_ids())], + ) + .unwrap(); + + let covered_schema = ArrowSchema::new(vec![ + vec_field, + ROW_ID_FIELD.clone(), + Field::new("extra", DataType::UInt32, true), + ]); + let covered = RecordBatch::try_new( + covered_schema.into(), + vec![ + Arc::new(make_fsl()), + Arc::new(row_ids()), + Arc::new(UInt32Array::from_iter_values((0..TOTAL).map(|v| v as u32))), + ], + ) + .unwrap(); + + let err = StorageBuilder::new("vec".to_owned(), pq.distance_type, pq, None) + .unwrap() + .build(vec![plain, covered]) + .expect_err("a batch with extra columns must be rejected, not silently projected"); + assert!(err.to_string().contains("extra"), "got: {err}"); + } + + #[tokio::test] + async fn test_new_drops_legacy_part_id_column() { + // Pre-#3606 `num_partitions=1` builds wrote `__ivf_part_id` into the + // index file. It must still be dropped on load (the legacy read shim), + // NOT preserved and misclassified as a covering column -- that would + // make search emit an extra column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks. + use crate::vector::PART_ID_COLUMN; + let codebook = Float32Array::from_iter_values((0..256 * DIM).map(|_| rand::random())); + let codebook = FixedSizeListArray::try_new_from_values(codebook, DIM as i32).unwrap(); + + let schema = ArrowSchema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new( + PQ_CODE_COLUMN, + DataType::FixedSizeList( + Field::new_list_field(DataType::UInt8, true).into(), + NUM_SUB_VECTORS as i32, + ), + true, + ), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + ]); + let row_ids = UInt64Array::from_iter_values((0..TOTAL).map(|v| v as u64)); + let codes = UInt8Array::from_iter_values((0..TOTAL * NUM_SUB_VECTORS).map(|v| v as u8)); + let codes = FixedSizeListArray::try_new_from_values(codes, NUM_SUB_VECTORS as i32).unwrap(); + let part_ids = UInt32Array::from_iter_values((0..TOTAL).map(|_| 0)); + let batch = RecordBatch::try_new( + schema.into(), + vec![Arc::new(row_ids), Arc::new(codes), Arc::new(part_ids)], + ) + .unwrap(); + + let storage = ProductQuantizationStorage::new( + codebook, + batch, + 8, + NUM_SUB_VECTORS, + DIM, + DistanceType::L2, + true, // codes already transposed + None, + ) + .unwrap(); + assert!( + storage.batch().column_by_name(PART_ID_COLUMN).is_none(), + "legacy __ivf_part_id must be dropped on load, not kept" + ); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + #[tokio::test] async fn test_distance_all() { let storage = create_pq_storage().await; @@ -1392,8 +1584,19 @@ mod tests { // Rewritten row i lands at offset i of frag 1. assert_eq!(*row_id, (1u64 << 32) | i as u64); } - assert_eq!(new_storage.batch.num_columns(), 2); + // The covering ("extra") column must survive remap, row-aligned to the + // surviving rows — otherwise PK/covering data is lost on compaction. + assert_eq!(new_storage.batch.num_columns(), 3); assert!(new_storage.batch.column_by_name(ROW_ID).is_some()); assert!(new_storage.batch.column_by_name(PQ_CODE_COLUMN).is_some()); + let extra = new_storage + .batch + .column_by_name("extra") + .expect("extra column should survive remap") + .as_primitive::(); + // Surviving rows are original indices 0..TOTAL/2, whose extra values are 0..TOTAL/2. + for i in 0..TOTAL / 2 { + assert_eq!(extra.value(i), i as u32); + } } } diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 151be945741..fb2be204dd5 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -373,6 +373,25 @@ impl DeepSizeOf for QueryScratchPool { /// /// It abstracts away the logic to compute the distance between vectors. /// +/// Indices of the covering ("included") columns in a vector-storage schema: every +/// field whose name is not one of `internal`. Each storage lists its own non-covering +/// column names (the row id, its quantization code columns, and any legacy bookkeeping +/// column it carries) and shares this filter, so covering detection stays a per-storage +/// data decision rather than duplicated iteration logic. See +/// [`VectorStore::covering_field_indices`] for why the set cannot be inferred generically. +pub(crate) fn covering_field_indices_excluding( + schema: &arrow_schema::Schema, + internal: &[&str], +) -> Vec { + schema + .fields() + .iter() + .enumerate() + .filter(|(_, f)| !internal.contains(&f.name().as_str())) + .map(|(i, _)| i) + .collect() +} + /// TODO: should we rename this to "VectorDistance"?; /// ///
@@ -410,6 +429,45 @@ pub trait VectorStore: Send + Sync + Sized + Clone { /// The storage implement will perform quantization if necessary. fn append_batch(&self, batch: RecordBatch, vector_column: &str) -> Result; + /// Field indices of the "included"/covering columns: the extra columns + /// stored alongside the row id and quantization code so a covered query can + /// skip the take from the base table. + /// + /// The default is empty (no covering). A storage that supports covering must + /// override this, since only it knows which of its columns are quantization + /// code columns versus included columns — inferring that generically by name + /// is unsafe (e.g. it would mistake RaBitQ's code columns for covering ones). + fn covering_field_indices(&self) -> Vec { + Vec::new() + } + + /// `[_rowid, ]` for the whole storage. Captured while a + /// partition is loaded during search so covering columns can be emitted with + /// the result without a separate take. Returns `None` if there are no + /// included columns (ordinary index — nothing to cover). + fn covering_batch(&self) -> Result> { + let included = self.covering_field_indices(); + if included.is_empty() { + return Ok(None); + } + let schema = self.schema().clone(); + let row_id_idx = schema.index_of(ROW_ID)?; + let mut indices = Vec::with_capacity(included.len() + 1); + indices.push(row_id_idx); + indices.extend(included); + // Project each batch BEFORE concatenating. Concatenating the full partition first + // would copy every quantization-code column only to discard it on the next line -- + // and those columns dominate the partition (HNSW partitions target 1 << 20 rows), so + // on a multi-probe query that copy is hundreds of MB of pure waste on the ANN hot + // path. Only `[_rowid, ]` is ever read from the result. + let projected: Vec = self + .to_batches()? + .map(|batch| batch.project(&indices)) + .collect::>()?; + let projected_schema = Arc::new(schema.project(&indices)?); + Ok(Some(concat_batches(&projected_schema, projected.iter())?)) + } + /// Create a [DistCalculator] to compute the distance between the query. /// /// Using dist calculator can be more efficient as it can pre-compute some @@ -480,7 +538,43 @@ impl StorageBuilder { } pub fn build(&self, batches: Vec) -> Result { - let mut batch = concat_batches(batches[0].schema_ref(), batches.iter())?; + // Batches can come from different sources (existing partitions loaded + // from disk vs freshly shuffled/split/reassigned data) that carry the + // same columns in a different order -- notably when the index has + // included/covering columns. `concat_batches` matches columns by + // position, so align every batch to the first batch's column order (by + // name). No-op when all batches already share a schema. + let schema = batches[0].schema(); + let aligned: Vec = batches + .iter() + .map(|b| { + if b.schema() == schema { + return Ok(b.clone()); + } + // `project_by_schema` reorders by name and errors on a missing + // column, but it cannot see extras -- and a batch with MORE + // columns must error rather than have them silently dropped + // (e.g. covering columns the index metadata still advertises). + // Equal count + every expected column present = same set. + if b.num_columns() != schema.fields().len() { + let names = |s: &arrow_schema::Schema| { + s.fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>() + .join(", ") + }; + return Err(Error::index(format!( + "mismatched columns while merging vector storage batches: \ + expected [{}], got [{}]", + names(&schema), + names(&b.schema()), + ))); + } + Ok(b.project_by_schema(schema.as_ref())?) + }) + .collect::>>()?; + let mut batch = concat_batches(&schema, aligned.iter())?; if batch.column_by_name(self.quantizer.column()).is_none() { let vectors = batch @@ -518,6 +612,13 @@ pub struct IvfQuantizationStorage { ivf: IvfModel, frag_reuse_index: Option>, + /// Lazily-computed covering schema (see [`Self::covering_schema`]). The schema is + /// fixed for the storage's lifetime and is consulted on every search, so it is + /// resolved once rather than per query. Constructing it builds an empty `Q::Storage`: + /// cheap for a current-format index (zero rows, and the codebook is cloned, not + /// decoded), but the legacy `codebook_tensor` path does decode, so this is not free + /// for every index. + covering_schema: std::sync::OnceLock>, } impl DeepSizeOf for IvfQuantizationStorage { @@ -589,6 +690,7 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), }) } @@ -620,6 +722,7 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), } } @@ -656,6 +759,31 @@ impl IvfQuantizationStorage { Arc::new(self.reader.schema().as_ref().into()) } + /// The `[_rowid, ]` schema for this index's covering + /// columns, or `None` if the index has no covering columns. Derived from an + /// empty storage instance (no I/O) so callers can emit the correct covered + /// output schema even when zero partitions are searched. Cached because it is queried + /// per search. The cache is best-effort under concurrency: racing cold callers may each + /// construct one and all but the first are discarded. That is deliberate -- the empty + /// construction is cheap on the current format, and `OnceLock::get_or_init` cannot + /// carry the `Result` this returns. + pub fn covering_schema(&self) -> Result> { + if let Some(cached) = self.covering_schema.get() { + return Ok(cached.clone()); + } + let arrow_schema = arrow_schema::Schema::from(self.reader.schema().as_ref()); + let empty = RecordBatch::new_empty(Arc::new(arrow_schema)); + // No remapper: the batch is empty, so there are no row ids to remap and the + // remapper cannot influence the schema this derives. Passing one would also be + // wrong now that the field holds a `dyn RowIdRemapper` -- the default + // `try_from_batch_with_remapper` rejects a remapper outright for every storage + // that does not override it, which would fail this call on an SQ/RQ/FLAT index + // whose dataset happens to carry a fragment-reuse index. + let storage = Q::Storage::try_from_batch(empty, self.metadata(), self.distance_type, None)?; + let computed = storage.covering_batch()?.map(|b| b.schema()); + Ok(self.covering_schema.get_or_init(|| computed).clone()) + } + /// Get the number of partitions in the storage. pub fn num_partitions(&self) -> usize { self.ivf.num_partitions() diff --git a/rust/lance-index/src/vector/transform.rs b/rust/lance-index/src/vector/transform.rs index 2faae7b9ce6..1ad89afd237 100644 --- a/rust/lance-index/src/vector/transform.rs +++ b/rust/lance-index/src/vector/transform.rs @@ -7,18 +7,18 @@ use std::fmt::Debug; use std::sync::Arc; -use arrow::datatypes::UInt64Type; -use arrow_array::UInt64Array; use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_array::{Array, ArrowPrimitiveType, RecordBatch, UInt32Array, cast::AsArray}; use arrow_schema::{DataType, Field, Schema}; use lance_arrow::RecordBatchExt; use num_traits::Float; -use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_core::{Error, Result}; use lance_linalg::kernels::normalize_fsl; use tracing::instrument; +use super::{CENTROID_DIST_COLUMN, PART_ID_COLUMN}; + /// Transform of a Vector Matrix. /// /// @@ -194,25 +194,48 @@ impl Transformer for Flatten { match arr.data_type() { DataType::FixedSizeList(_, _) => Ok(batch.clone()), DataType::List(_) => { - let row_ids = batch[ROW_ID].as_primitive::(); let vectors = arr.as_list::(); - - let row_ids = row_ids.values().iter().zip(vectors.iter()).flat_map( - |(row_id, multivector)| { - std::iter::repeat_n( - *row_id, - multivector.map(|multivec| multivec.len()).unwrap_or(0), - ) - }, - ); - let row_ids = UInt64Array::from_iter_values(row_ids); - let vectors = vectors.values().as_fixed_size_list().clone(); - let schema = Arc::new(Schema::new(vec![ - ROW_ID_FIELD.clone(), - Field::new(self.column.as_str(), vectors.data_type().clone(), true), - ])); - let batch = - RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(vectors)])?; + // Each source row expands into one output row per vector in its list. + // Replicate the row id AND every other column (e.g. covering / "included" + // columns) across the expansion via a gather; dropping them here would + // silently lose covering data on the multivector build/split path. + // + // The IVF partition-assignment columns are the exception and must be + // DROPPED, not replicated: they describe the source *row*, and after the + // expansion each sub-vector needs its own nearest centroid. Carrying them + // through makes `PartitionTransformer` see its output column already + // present and early-return, silently collapsing every sub-vector of a row + // into that row's single partition -- no error, just degraded recall. + let take_indices = + UInt32Array::from_iter_values((0..vectors.len()).flat_map(|i| { + let n = if vectors.is_valid(i) { + vectors.value(i).len() + } else { + 0 + }; + std::iter::repeat_n(i as u32, n) + })); + let flat_vectors = vectors.values().as_fixed_size_list().clone(); + let mut fields: Vec> = Vec::with_capacity(batch.num_columns()); + let mut columns: Vec = + Vec::with_capacity(batch.num_columns()); + for (i, field) in batch.schema().fields().iter().enumerate() { + if field.name() == &self.column { + fields.push(Arc::new(Field::new( + self.column.as_str(), + flat_vectors.data_type().clone(), + true, + ))); + columns.push(Arc::new(flat_vectors.clone())); + } else if field.name() == PART_ID_COLUMN || field.name() == CENTROID_DIST_COLUMN + { + continue; + } else { + fields.push(field.clone()); + columns.push(arrow::compute::take(batch.column(i), &take_indices, None)?); + } + } + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; Ok(batch) } _ => Err(Error::index(format!( @@ -236,6 +259,114 @@ mod tests { use lance_arrow::*; use lance_linalg::distance::L2; + #[test] + fn test_flatten_preserves_extra_columns_for_multivector() { + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::UInt64Type; + use arrow_array::{ListArray, UInt64Array}; + use lance_core::{ROW_ID, ROW_ID_FIELD}; + + // row 10 has 2 vectors, row 20 has 1 vector; a covering column "meta" must be + // replicated across each row's expanded vectors, not dropped. + let values = Float32Array::from(vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]); + let fsl = FixedSizeListArray::try_new_from_values(values, 2).unwrap(); + let offsets = OffsetBuffer::new(vec![0i32, 2, 3].into()); + let item = Arc::new(Field::new("item", fsl.data_type().clone(), true)); + let list = ListArray::new(item, offsets, Arc::new(fsl), None); + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("vector", list.data_type().clone(), true), + Field::new("meta", DataType::Int32, true), + ])), + vec![ + Arc::new(UInt64Array::from(vec![10u64, 20])), + Arc::new(list), + Arc::new(Int32Array::from(vec![100, 200])), + ], + ) + .unwrap(); + + let out = Flatten::new("vector").transform(&batch).unwrap(); + + assert_eq!(out.num_rows(), 3, "2 + 1 vectors expand to 3 rows"); + assert_eq!( + out[ROW_ID].as_primitive::().values(), + &[10, 10, 20] + ); + let meta = out + .column_by_name("meta") + .expect("covering column must survive Flatten") + .as_primitive::(); + assert_eq!( + meta.values(), + &[100, 100, 200], + "covering column replicated per expanded vector" + ); + assert!(matches!( + out.column_by_name("vector").unwrap().data_type(), + DataType::FixedSizeList(_, 2) + )); + } + + /// The IVF partition-assignment columns describe the source *row*. Replicating them + /// across the multivector expansion makes `PartitionTransformer` see its output column + /// already present and early-return, so every sub-vector inherits its row's single + /// partition instead of being assigned to its own nearest centroid -- silent recall + /// loss with no error. They must be dropped even though covering columns are kept. + #[test] + fn test_flatten_drops_partition_columns_for_multivector() { + use arrow::buffer::OffsetBuffer; + use arrow_array::{ListArray, UInt32Array as U32, UInt64Array}; + use lance_core::ROW_ID_FIELD; + + let values = Float32Array::from(vec![1.0, 1.0, 2.0, 2.0, 3.0, 3.0]); + let fsl = FixedSizeListArray::try_new_from_values(values, 2).unwrap(); + let offsets = OffsetBuffer::new(vec![0i32, 2, 3].into()); + let item = Arc::new(Field::new("item", fsl.data_type().clone(), true)); + let list = ListArray::new(item, offsets, Arc::new(fsl), None); + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + ROW_ID_FIELD.clone(), + Field::new("vector", list.data_type().clone(), true), + Field::new("meta", DataType::Int32, true), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + Field::new(CENTROID_DIST_COLUMN, DataType::Float32, true), + ])), + vec![ + Arc::new(UInt64Array::from(vec![10u64, 20])), + Arc::new(list), + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(U32::from(vec![7u32, 9])), + Arc::new(Float32Array::from(vec![0.5f32, 0.25])), + ], + ) + .unwrap(); + + let out = Flatten::new("vector").transform(&batch).unwrap(); + + assert_eq!(out.num_rows(), 3); + assert!( + out.column_by_name(PART_ID_COLUMN).is_none(), + "a row's partition assignment must not survive the multivector expansion, or \ + PartitionTransformer early-returns and every sub-vector shares one partition" + ); + assert!( + out.column_by_name(CENTROID_DIST_COLUMN).is_none(), + "the row's centroid distance is likewise meaningless per sub-vector" + ); + // ...while genuine covering columns are still replicated. + assert_eq!( + out.column_by_name("meta") + .expect("covering column must survive") + .as_primitive::() + .values(), + &[100, 100, 200] + ); + } + #[tokio::test] async fn test_normalize_transformer_f32() { let data = Float32Array::from_iter_values([1.0, 1.0, 2.0, 2.0].into_iter()); diff --git a/rust/lance/src/dataset/index.rs b/rust/lance/src/dataset/index.rs index a3e0917d432..ff46d27adae 100644 --- a/rust/lance/src/dataset/index.rs +++ b/rust/lance/src/dataset/index.rs @@ -421,79 +421,6 @@ mod tests { assert_ne!(remapped[0].new_id, unaffected_segment_id); } - /// A covered index must be withdrawn from remapping rather than remapped. - /// No index type carries the declared payload through a remap, so a - /// replacement would republish a covering claim its storage does not back. - /// Withdrawal also means the legacy `index_details: None` migration path is - /// never reached for such an index, so it cannot panic there either. - #[tokio::test] - async fn test_remapper_migration_path_withdraws_covered_index() { - let reader = lance_datagen::gen_batch() - .col("a", array::step::()) - .col("b", array::step::()) - .into_reader_rows(RowCount::from(20), BatchCount::from(1)); - let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); - dataset - .create_index( - &["a"], - IndexType::BTree, - None, - &ScalarIndexParams::for_builtin(BuiltinIndexType::BTree), - false, - ) - .await - .unwrap(); - - let a_id = dataset.schema().field("a").unwrap().id; - let b_id = dataset.schema().field("b").unwrap().id; - let current = dataset.load_indices().await.unwrap(); - let mut legacy_covered = current[0].clone(); - legacy_covered.fields = vec![a_id, b_id]; - legacy_covered.covering_fields = vec![b_id]; - // Force the legacy migration path this fix touches. - legacy_covered.index_details = None; - - let transaction = Transaction::new( - dataset.manifest.version, - Operation::CreateIndex { - new_indices: vec![legacy_covered], - removed_indices: current.to_vec(), - }, - None, - ); - dataset - .apply_commit(transaction, &Default::default(), &Default::default()) - .await - .unwrap(); - - let index_uuid = dataset.load_indices().await.unwrap()[0].uuid; - - // Fully delete every row so `remap_index` returns `RemapResult::Keep`, - // landing in the `index_details: None` migration branch this fix touches. - let remap_to_empty = (0..dataset.count_all_rows().await.unwrap()) - .map(|i| (i as u64, None)) - .collect::>(); - let remapper = DatasetIndexRemapperOptions::default() - .create_remapper(&dataset) - .await - .unwrap() - .expect("a real index should require a remapper"); - let remapped = remapper - .remap_indices(RowAddrRemap::direct(remap_to_empty), &[0]) - .await - .unwrap(); - - // Withdrawn, not remapped: no index type carries the declared payload - // through a remap, so producing a replacement would republish a covering - // claim its storage does not back. The original entry stays in the - // manifest and simply stops covering the rewritten fragments. - assert!( - remapped.is_empty(), - "a covered index must be withdrawn from remapping, got {remapped:?}" - ); - let _ = index_uuid; - } - /// The same migration path must reject -- cleanly, not with a panic -- a /// malformed `covering_fields` longer than `fields`. /// diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 81f01882e55..5f8b4f06645 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -6994,6 +6994,7 @@ mod tests { version: crate::index::vector::IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), }, false, ) @@ -7129,6 +7130,7 @@ mod tests { version: crate::index::vector::IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), }, false, ) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 87f2d83a281..21c93c051cb 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5649,11 +5649,32 @@ impl Scanner { knn_node = self.take(knn_node, vector_projection)?; } + // The index search now emits the index's covering ("included") columns, so every flat + // path below must produce the same columns before it is unioned with `knn_node`; + // otherwise projecting the flat output to `knn_node.schema()` panics on the missing column. + let covering_columns: Vec = indexed_segments + .first() + .map(|s| s.covering_fields.as_slice()) + .unwrap_or(&[]) + .iter() + .filter_map(|id| { + self.dataset + .schema() + .field_by_id(*id) + .map(|field| field.name.clone()) + }) + .collect(); + let mut columns = vec![q.column.clone()]; if let Some(expr) = filter_plan.full_expr.as_ref() { let filter_columns = Planner::column_names_in_expr(expr); columns.extend(filter_columns); } + for name in &covering_columns { + if !columns.contains(name) { + columns.push(name.clone()); + } + } // Collect flat-path plans; union order matches original (flat before ANN) so test snapshots // and downstream plan analyses remain stable. @@ -5708,6 +5729,9 @@ impl Scanner { let filter_columns = Planner::column_names_in_expr(expr); take_proj = take_proj.union_columns(filter_columns, OnMissing::Error)?; } + if !covering_columns.is_empty() { + take_proj = take_proj.union_columns(covering_columns.clone(), OnMissing::Error)?; + } let mut stale_node = self.stale_rows_take(stale_rows, take_proj).await?; if let Some(expr) = filter_plan.full_expr.as_ref() { stale_node = Arc::new(LanceFilterExec::try_new(expr.clone(), stale_node)?); @@ -15032,6 +15056,7 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") version: crate::index::vector::IndexFileVersion::Legacy, skip_transpose: false, runtime_hints: Default::default(), + covering_columns: Default::default(), }, false, ) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 19673e83ae1..e4eddb6a321 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -13,7 +13,7 @@ use crate::dataset::WriteDestination; use crate::dataset::builder::DatasetBuilder; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; -use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; +use crate::dataset::transaction::{DataReplacementGroup, Operation}; use crate::index::vector::VectorIndexParams; use crate::session::Session; use crate::utils::test::covering; @@ -21,7 +21,7 @@ use crate::{Dataset, Error, Result}; use lance_arrow::FixedSizeListArrayExt; use crate::dataset::write::{WriteMode, WriteParams}; -use crate::index::DatasetIndexExt; +use crate::index::{CreateIndexBuilder, DatasetIndexExt}; use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; use arrow::datatypes::UInt64Type; use arrow_array::RecordBatch; @@ -56,6 +56,8 @@ use lance_index::{IndexType, scalar::ScalarIndexParams, vector::DIST_COL}; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; use lance_linalg::distance::MetricType; +use lance_table::format::DataFile; +use object_store::path::Path; use datafusion::common::{assert_contains, assert_not_contains}; use futures::{StreamExt, TryStreamExt}; @@ -64,7 +66,7 @@ use lance_arrow::json::ARROW_JSON_EXT_NAME; use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; use lance_table::format::BasePath; use lance_testing::datagen::generate_random_array; -use rand::Rng; +use rand::{Rng, SeedableRng, rngs::StdRng}; use rstest::rstest; #[rstest] @@ -225,6 +227,125 @@ async fn test_covered_vector_column_is_not_selected_for_ann() { ); } +/// The entire point of a covered IVF_PQ index: a query whose projection is +/// satisfied by the keyed vector column's own `_distance`/`_rowid` plus the +/// index's declared `covering_fields` must be answered from the index alone, +/// with no take against the base table. A query that also needs a column the +/// index does not carry must still take. +/// +/// Two traps established empirically on this project make a broken +/// take-elision path look correct anyway: +/// - `num_partitions = 1` never reaches the probe path at all. +/// - Even 4 partitions is not enough with uniform-random vectors: the +/// early-pruning heuristic (`early_pruning` in `io/exec/knn.rs`) still +/// searches every partition, so `late_search`'s split is never +/// exercised. The data must have well-separated clusters, and +/// `early_pruning` must stay off. +/// So this test uses 4 well-separated clusters and asserts on the plan the +/// scanner actually built (`explain_plan`), never on query results -- a take +/// can run and simply return the same values. +#[tokio::test] +async fn test_covered_ann_query_elides_base_table_take() { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let test_uri = TempStrDir::default(); + + // Well-separated clusters (centers 1000 apart along dim 0), not + // uniform-random vectors -- see the module doc above. + let mut rng = StdRng::seed_from_u64(42); + let mut payload = Vec::with_capacity(TOTAL); + let mut extra = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let row_id = (cluster * ROWS_PER_CLUSTER + row) as i32; + payload.push(row_id); + extra.push(row_id * 2); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + + let vectors: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vector", vectors.data_type().clone(), false), + ArrowField::new("payload", DataType::Int32, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + vectors, + Arc::new(Int32Array::from(payload)), + Arc::new(Int32Array::from(extra)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); + + // 4 partitions, matching the cluster count: the covering columns' + // reachability doesn't hinge on which internal probe path runs, but a + // covered-ANN test with too few effectively-reachable partitions has + // already once passed against broken code on this project (see doc + // above), so this follows the same defensive recipe regardless. + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + // Query near the last cluster's center. + let mut q_values = vec![0.0f32; DIMS]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let query = Float32Array::from(q_values); + + // Covered case: the projection (`payload`) is entirely satisfied by the + // index's own output columns -- the plan must contain no base-table take. + // A base-table take on this (stable/v2) storage format is a + // `FilteredReadExec`, displayed as `LanceRead: ...`; checking for the + // literal substring "Take" would be a no-op here, since that name is + // only ever used by the legacy (v1) `TakeExec`. + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "a covered projection ['payload'] must skip the base-table take entirely:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("payload").is_some()); + + // Uncovered case: `extra` is not a declared covering column, so the + // scanner must take it from the base table. `source=stream` is the + // marker for a row-id-driven take (as opposed to a full-table scan), + // which is what a base-table take against ANN results looks like. + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["extra"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead") && plan.contains("source=stream"), + "a projection needing the uncovered column 'extra' must take from the base table:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("extra").is_some()); +} + /// A filtered `describe_indices` must still find a covered index by its keyed /// column. The matcher compares the caller's resolved field slice against the one /// column named by `for_column`, so a caller that passes all of `index.fields` -- @@ -273,85 +394,23 @@ async fn test_describe_indices_filters_a_covered_index_by_its_keyed_field() { ); } -/// An unfiltered `optimize_indices()` is a table-wide request, so a covered index -/// must not abort it -- that would block optimization of every other index on the -/// table over one this build merely cannot rebuild. It is skipped with a warning. -/// Only a caller that names the covered index gets an error (see -/// `test_optimize_indices_rejects_a_covered_index`). +/// The scalar counterpart to the covered-vector optimize case above. Before a producer +/// existed optimize skipped a covered group outright; the producer makes covered rebuilds +/// real, so what needs pinning is that rebuilding a covered *scalar* group preserves its +/// covering declaration rather than silently dropping it. A carried column is inert to +/// every scalar read path, so a lost declaration would stay invisible until a covered +/// planner consulted it. /// -/// Both the current and the stale case are covered: erroring on the stale one -/// aborts the loop before the replacements accumulated for the other groups are -/// committed, leaving an unrelated index stale too. -#[rstest] -#[case::current(false)] -#[case::stale(true)] -#[tokio::test] -async fn test_optimize_skips_a_covered_index_without_blocking_others(#[case] stale: bool) { - let test_uri = TempStrDir::default(); - let mut dataset = covering::write_vector_payload_dataset(&test_uri).await; - - // A covered vector index, and a plain scalar index that optimize may touch. - covering::create_ivf_pq_index(&mut dataset, "vec").await; - covering::create_btree_index(&mut dataset, "payload", Some("payload_idx")).await; - - let (_, payload_id) = covering::declare_covering(&mut dataset, "vec", "payload").await; - let covered_uuid = dataset - .load_indices() - .await - .unwrap() - .iter() - .find(|idx| !idx.covering_fields.is_empty()) - .expect("the covered index should exist") - .uuid; - - if stale { - // Now *both* groups have an unindexed fragment, so the covered group - // would genuinely be rebuilt -- this is where the refusal used to fire. - covering::append_vector_payload_rows(&mut dataset, 256).await; - } - - dataset - .optimize_indices(&OptimizeOptions::default()) - .await - .expect("a covered index must not abort an unfiltered optimize"); - - let after = dataset.load_indices().await.unwrap(); - assert!( - after - .iter() - .any(|idx| idx.uuid == covered_uuid && idx.covering_fields == vec![payload_id]), - "the covered index must be left exactly as it was" - ); - - if stale { - let payload_idx = after - .iter() - .filter(|idx| idx.name == "payload_idx") - .filter_map(|idx| idx.fragment_bitmap.as_ref()) - .fold(roaring::RoaringBitmap::new(), |mut acc, bitmap| { - acc |= bitmap; - acc - }); - assert!( - payload_idx.contains(1), - "the unrelated index must still have been optimized onto the new fragment, got {payload_idx:?}" - ); - } else { - assert_eq!(after.len(), 2, "both indices must survive"); - } -} - -/// The same skip, for a *scalar* covered index. The rule does not branch on index -/// type, but scalar groups take their own no-work path a few lines below the -/// covering gate, so a covered scalar index reaching that gate first is worth -/// pinning separately from the vector case above. +/// The append is what gives this test teeth: with no unindexed rows the covered group has +/// no work, and the assertions would hold with the rebuild never happening. /// -/// The append is what gives this test teeth. Without it the covered group has no -/// work either way, so the scalar no-work path below the gate produces the same -/// observable outcome as the gate itself and the test passes with the gate -/// removed entirely. +/// The reload right before `optimize_indices` is load-bearing. Without it optimize +/// intermittently commits a manifest missing one of the two groups. That race needs no +/// covering to reproduce -- two plain btree indexes over an append are enough -- so it is +/// not this branch's to fix, and the reload keeps this test measuring covering rather than +/// that race. #[tokio::test] -async fn test_optimize_skips_a_stale_covered_scalar_index() { +async fn test_optimize_rebuild_preserves_a_covered_scalar_declaration() { let test_uri = TempStrDir::default(); let mut dataset = covering::write_three_int_column_dataset(&test_uri).await; covering::create_btree_index(&mut dataset, "a", None).await; @@ -367,9 +426,10 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { .expect("the covered index should exist") .uuid; - // Both scalar groups now have an unindexed fragment, so the covered one - // would genuinely be rebuilt if the gate did not skip it first. + // Both scalar groups now have an unindexed fragment, so the covered one is genuinely + // rebuilt rather than reaching a no-work path. covering::append_three_int_column_rows(&mut dataset, 64).await; + dataset.load_indices().await.unwrap(); dataset .optimize_indices(&OptimizeOptions::default()) @@ -377,14 +437,36 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { .expect("a stale covered scalar index must not abort optimize"); let after = dataset.load_indices().await.unwrap(); + let covered = after + .iter() + .filter(|idx| !idx.covering_fields.is_empty()) + .collect::>(); + assert_eq!( + covered.len(), + 1, + "exactly one index must still declare covering, got {covered:?}" + ); + assert_eq!( + covered[0].covering_fields, + vec![carried_id], + "the rebuild must carry the covering declaration forward, not drop it" + ); + assert_ne!( + covered[0].uuid, covered_uuid, + "precondition: the covered group must actually have been rebuilt, or this test \ + would pass with the rebuild path never entered" + ); + let covered_coverage = covered[0] + .fragment_bitmap + .as_ref() + .expect("a rebuilt index records its coverage"); assert!( - after - .iter() - .any(|idx| idx.uuid == covered_uuid && idx.covering_fields == vec![carried_id]), - "the covered scalar index must be left exactly as it was, not rebuilt" + covered_coverage.contains(1), + "the rebuilt covered index must span the appended fragment, got {covered_coverage:?}" ); - // The unrelated scalar index was still maintained, so the skip is scoped to - // the covered group rather than aborting the loop. + + // The unrelated scalar index was still maintained, so the rebuild is scoped to the + // covered group rather than aborting the loop. let b_id = dataset.schema().field_id("b").unwrap(); let b_coverage = after .iter() @@ -400,60 +482,6 @@ async fn test_optimize_skips_a_stale_covered_scalar_index() { ); } -/// A caller that names the covered index asked for it specifically, so the -/// refusal is loud rather than a skip. -#[tokio::test] -async fn test_optimize_indices_rejects_a_covered_index() { - let test_uri = TempStrDir::default(); - let mut dataset = covering::write_vector_payload_dataset(&test_uri).await; - covering::create_ivf_pq_index(&mut dataset, "vec").await; - - // Nothing writes carried values, so the storage does not contain `payload` - // -- which is exactly why optimize must refuse: it rebuilds from a scan - // projecting the keyed field and `_rowid` only, and would republish the - // declaration on a segment that still has no payload. - let (vec_id, payload_id) = covering::declare_covering(&mut dataset, "vec", "payload").await; - - // Append AFTER declaring covering, so the group really would be rebuilt. - // Without it this would assert the refusal against an index optimize had no - // work for, and would keep passing if the refusal moved behind a no-work - // check. - covering::append_vector_payload_rows(&mut dataset, 256).await; - - let before = dataset.load_indices().await.unwrap(); - let before_uuid = before[0].uuid; - let covered_name = before[0].name.clone(); - - // Name the covered index: the refusal is reserved for a caller that targeted - // it. An unfiltered call skips it instead, which - // `test_optimize_skips_a_covered_index_without_blocking_others` covers. - let err = dataset - .optimize_indices(&OptimizeOptions::default().index_names(vec![covered_name])) - .await - .expect_err("optimizing a targeted covered index must be refused"); - assert!( - err.to_string().contains("declares covering fields"), - "unexpected message: {err}" - ); - - // Refused, not partially applied: the index is exactly as it was. - let after = dataset.load_indices().await.unwrap(); - assert_eq!(after.len(), 1); - assert_eq!( - after[0].uuid, before_uuid, - "a refused optimize must not replace the index" - ); - assert_eq!(after[0].covering_fields, vec![payload_id]); - assert_eq!(after[0].fields, vec![vec_id, payload_id]); - - // The appended data above is unindexed, so this really is a case optimize - // would otherwise have merged -- the refusal is not the no-op path. - assert!( - after.iter().all(|idx| !idx.covering_fields.is_empty()), - "precondition: the only index is still the covered one" - ); -} - #[rstest] #[tokio::test] async fn test_create_scalar_index( @@ -7448,17 +7476,91 @@ async fn test_load_segment_params_full_fidelity() { assert_eq!(&read, opened.params()); } -/// Compaction of a covered index must succeed. `remap_index` rejected any -/// index with more than one field, so this failed outright -- a covered -/// dataset could be created but never compacted. -#[tokio::test] -async fn test_compaction_withdraws_a_covered_index_without_failing() { - use crate::dataset::optimize::{CompactionOptions, compact_files}; +/// Write `values` as a standalone single-field `payload` data file (field `payload_id`) +/// and return the resulting [`DataFile`], ready to use in an `Operation::DataReplacement` +/// group. Kept in its own file from `vec` throughout this test so replacing it is a +/// straightforward single-field swap, never a split of a shared file. +async fn write_payload_data_file( + dataset: &Dataset, + name: &str, + payload_id: i32, + values: Int32Array, +) -> DataFile { + let payload_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "payload", + DataType::Int32, + true, + )])); + let path = format!("{name}.lance"); + let object_writer = dataset + .object_store + .create(&Path::from(format!("data/{path}"))) + .await + .unwrap(); + let mut writer = lance_file::versions::v2_1::create_writer( + object_writer, + payload_schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + let write_batch = RecordBatch::try_new(payload_schema, vec![Arc::new(values)]).unwrap(); + writer.write_batch(&write_batch).await.unwrap(); + writer.finish().await.unwrap(); - let test_uri = TempStrDir::default(); + let (major, minor) = LanceFileVersion::Stable.resolve().to_data_file_numbers(); + DataFile { + path, + fields: Arc::from([payload_id]), + column_indices: Arc::from([0]), + file_major_version: major, + file_minor_version: minor, + file_size_bytes: CachedFileSize::unknown(), + base_id: None, + } +} + +/// `prune_stale_segment_coverage` must judge staleness from every field a segment +/// declares -- keyed and carried alike -- not just the keyed field's subtree. +/// +/// Before the fix, only the keyed (`vec`) field fed the staleness check: a fragment +/// whose *carried* (`payload`) column was physically rewritten between when the +/// segment was built and when it is committed stayed marked "covered" regardless, so +/// the committed index would go on serving that fragment's now-stale carried values. +/// `dataset/overlay.rs`'s per-read staleness check already considers every declared +/// field; this is the build-time counterpart disagreeing with it. +#[tokio::test] +async fn test_commit_existing_index_segments_prunes_fragment_with_rewritten_carried_column() { + // An in-memory store, not `TempStrDir`: `write_payload_data_file` writes bare + // relative paths (matching this crate's other `DataReplacement` tests), which for + // a `file://` store would resolve against the filesystem root rather than the + // dataset's own directory. + let test_uri = "memory://"; let dimension = 16; - let schema = Arc::new(ArrowSchema::new(vec![ + let vec_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + dimension, + ), + false, + )])); + let vectors = Arc::new( + ::try_new_from_values( + generate_random_array(512 * dimension as usize), + dimension, + ) + .unwrap(), + ); + let batch = RecordBatch::try_new(vec_schema.clone(), vec![vectors]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], vec_schema.clone()); + let dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + // Add `payload` as an all-null column. It stays fileless until populated below, so + // it never shares a data file with `vec` -- a later replacement of just its file is + // then a straightforward single-field swap, not a split of a multi-field file (which + // `DataReplacement` does not support). + let extended_schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new( "vec", DataType::FixedSizeList( @@ -7467,81 +7569,105 @@ async fn test_compaction_withdraws_a_covered_index_without_failing() { ), false, ), - ArrowField::new("payload", DataType::Int32, false), + ArrowField::new("payload", DataType::Int32, true), ])); + let fragment = dataset.get_fragments().pop().unwrap().metadata; + let read_version = dataset.manifest.version; + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::Merge { + fragments: vec![fragment], + schema: extended_schema.as_ref().try_into().unwrap(), + preserves_nullability: true, + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + let payload_id = dataset.schema().field_id("payload").unwrap(); - let make_batch = |offset: i32| { - let vectors = Arc::new( - ::try_new_from_values( - generate_random_array(256 * dimension as usize), - dimension, - ) - .unwrap(), - ); - let payload = Arc::new(Int32Array::from_iter_values(offset..offset + 256)); - RecordBatch::try_new(schema.clone(), vec![vectors, payload]).unwrap() - }; - - // Two fragments, so compaction has something to compact. - let reader = RecordBatchIterator::new(vec![Ok(make_batch(0))], schema.clone()); - let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(make_batch(256))], schema.clone()); - dataset.append(reader, None).await.unwrap(); + // Give `payload` its first real values, in a data file of its own. + let payload_v1 = write_payload_data_file( + &dataset, + "payload_v1", + payload_id, + Int32Array::from_iter_values(0..512), + ) + .await; + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, payload_v1)], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); - let params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 50); - dataset - .create_index(&["vec"], IndexType::Vector, None, ¶ms, true) + // Build a covered segment (Task 1's `covering_columns`) while `payload` still holds + // its first values, but do not commit it yet. This mirrors a distributed build, + // whose segment metadata records the dataset version it was built from. + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 50); + params.covering_columns(vec!["payload".to_string()]); + let built = CreateIndexBuilder::new(&mut dataset, &["vec"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .execute_uncommitted() .await .unwrap(); let vec_id = dataset.schema().field_id("vec").unwrap(); - let payload_id = dataset.schema().field_id("payload").unwrap(); - let current = dataset.load_indices().await.unwrap(); - let mut covered = current[0].clone(); - covered.fields = vec![vec_id, payload_id]; - covered.covering_fields = vec![payload_id]; - - let transaction = Transaction::new( - dataset.manifest.version, - Operation::CreateIndex { - new_indices: vec![covered], - removed_indices: current.to_vec(), + assert_eq!(built.fields, vec![vec_id, payload_id]); + assert_eq!(built.covering_fields, vec![payload_id]); + assert!( + built.fragment_bitmap.as_ref().unwrap().contains(0), + "sanity: the segment must cover fragment 0 before any rewrite" + ); + + // Rewrite the *carried* `payload` column's data file -- `vec`'s file is untouched -- + // advancing the dataset past the version the segment above was built at. + let payload_v2 = write_payload_data_file( + &dataset, + "payload_v2", + payload_id, + Int32Array::from_iter_values(1000..1512), + ) + .await; + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, payload_v2)], }, + Some(read_version), None, - ); + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + dataset - .apply_commit(transaction, &Default::default(), &Default::default()) + .commit_existing_index_segments("vec_idx", "vec", vec![built]) .await .unwrap(); - let fragments_before: Vec = dataset.fragments().iter().map(|f| f.id).collect(); - assert!( - fragments_before.len() > 1, - "precondition: there must be something to compact" - ); - - // Compaction of the table must not be blocked by an index it cannot remap. - compact_files(&mut dataset, CompactionOptions::default(), None) - .await - .expect("compaction of a covered index must succeed"); - - let fragments_after: Vec = dataset.fragments().iter().map(|f| f.id).collect(); - assert_ne!( - fragments_after, fragments_before, - "compaction rewrote nothing, so the remap path never ran" - ); - - // The entry survives untouched -- withdrawal skips remapping rather than - // deleting metadata -- but it now covers none of the rewritten fragments, so - // no query can be answered from a payload the storage never held. - let after = dataset.load_indices().await.unwrap(); - assert_eq!(after.len(), 1); - assert_eq!(after[0].covering_fields, vec![payload_id]); - let live: roaring::RoaringBitmap = dataset.fragments().iter().map(|f| f.id as u32).collect(); - let effective = after[0].effective_fragment_bitmap(&live); + let committed = dataset.load_indices_by_name("vec_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + let coverage = committed[0].fragment_bitmap.as_ref().unwrap(); assert!( - effective.is_none_or(|bitmap| bitmap.is_empty()), - "a withdrawn covered index must stop covering fragments, got {:?}", - after[0].fragment_bitmap + !coverage.contains(0), + "fragment 0's carried `payload` column was rewritten after the segment was built \ + at dataset version {}; its coverage must be pruned, got {coverage:?}", + committed[0].dataset_version, ); } diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index b5803624520..6cc7affee02 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -163,7 +163,16 @@ struct PhysicalDataFileIdentity<'a> { base_id: Option, base_binding: PhysicalBaseBinding<'a>, path: &'a str, - fields: &'a [i32], + // Deliberately NOT the file's `fields` list. This identity is built per field, but + // `fields` is file-wide, so a tombstone of any *other* field in a shared file (an + // unrelated `RewriteColumns` turns its id negative) would change every field's + // identity and prune coverage for indexes whose data did not move. A tombstone of + // *this* field is still caught: its id stops matching and it drops out of the map + // `fragment_field_files` builds. `column_indices` covers an actual layout change. + // + // Upstream cannot hit this without a covering producer -- with no payload written, + // two indexes never share a file with different field sets -- which is why the guard + // held until this commit introduced one. column_indices: &'a [i32], file_major_version: u32, file_minor_version: u32, @@ -185,7 +194,6 @@ impl<'a> PhysicalDataFileIdentity<'a> { base_id: file.base_id, base_binding, path: &file.path, - fields: file.fields.as_ref(), column_indices: file.column_indices.as_ref(), file_major_version: file.file_major_version, file_minor_version: file.file_minor_version, @@ -1170,34 +1178,21 @@ pub(crate) async fn remap_index( // withdrawal below would otherwise swallow it as an ordinary covered index. matched.validate_covering_fields()?; - // A covered index cannot survive a remap yet. Nothing writes the carried - // values in the first place, and each index type's `remap` rewrites only the - // schema that type knows about, so the remapped segment would still declare - // payload its storage no longer holds. Withdraw it instead, exactly as an - // index whose type reports `can_remap() == false` is withdrawn below: an - // absent index costs a fallback scan, whereas a surviving false declaration - // is answered from data that is not there. Erroring here would instead block - // compaction of the whole table. - if !matched.covering_fields.is_empty() { - log::warn!( - "Index '{}' declares covering fields {:?}, which no index builder \ - writes or preserves yet. Index will be dropped during compaction \ - and must be rebuilt.", - matched.name, - matched.covering_fields, - ); - return Ok(RemapResult::Drop); - } - - // With covering withdrawn above, `fields` is entirely keyed, so more than one - // entry means a genuinely composite index. - if matched.fields.len() > 1 { + // Remap handles a single keyed field. Carried fields are irrelevant to the + // remap itself: the index is opened by its keyed field, and each storage + // format's `remap` preserves its covering columns row-aligned (the format + // commit dropped covered indexes here because nothing wrote carried values + // yet; the builders in this commit do). Only a genuinely composite index is + // rejected -- `keyed_fields` fails closed on malformed metadata, but the + // validation above already rejected that. + if matched.keyed_fields().len() > 1 { return Err(Error::index(format!( - "Remapping index '{}' is not supported: it has {} keyed fields {:?}; \ - only one keyed field is supported", + "Remapping index '{}' is not supported: it has {} keyed fields {:?} \ + (carried fields {:?}); only one keyed field is supported", matched.name, - matched.fields.len(), + matched.keyed_fields().len(), matched.fields, + matched.covering_fields, ))); } @@ -2311,42 +2306,6 @@ impl DatasetIndexExt for Dataset { continue; } - // Optimizing a covered index would republish its declaration on a - // segment rebuilt without the carried values: `scan_vector_fragments` - // projects the keyed field and `_rowid` only, and the scalar merges - // reconstruct value plus row id. - // - // What decides is the caller's intent, not whether this group is - // stale. An unfiltered `optimize_indices()` is a table-wide - // maintenance request, and erroring aborts the loop before the - // replacements accumulated for the other groups are committed -- so - // one index this build cannot rebuild would leave every other index - // on the table stale. Skip it with a warning instead. - // - // A caller that listed this index in `index_names` asked for it - // specifically, so refuse out loud. The loop is already filtered by - // that list, so reaching here with it set means this group was named. - if let Some(covered) = deltas - .iter() - .find(|index| !index.covering_fields.is_empty()) - { - if options.index_names.is_none() { - log::warn!( - "Skipping index '{}': it declares covering fields {:?}, \ - which no index builder writes or preserves yet.", - covered.name, - covered.covering_fields, - ); - continue; - } - return Err(Error::index(format!( - "Optimizing index '{}' is not supported: it declares \ - covering fields {:?}, which no index builder writes or \ - preserves yet", - covered.name, covered.covering_fields, - ))); - } - // Optimizing a name means replacing its segments with one that // covers their union, which this build cannot compute when it // cannot read one of them: the merged segment would overlap the @@ -8953,54 +8912,98 @@ mod tests { /// A segment's carried columns can go stale independently of its keyed column, /// so the staleness check has to walk every entry of `fields`, not just the - /// keyed subtree. Here the carried column did not exist when the segment was - /// built, so the segment cannot be carrying its values and its coverage of that - /// fragment must be pruned. Walking only the keyed subtree sees no change and - /// leaves the fragment covered, which is the bug. + /// keyed subtree. Here an in-place update rewrites the carried column's data + /// file after the segment was built, so the segment's stored copy is stale and + /// its coverage of that fragment must be pruned. Walking only the keyed subtree + /// sees no change (the vector file is untouched) and leaves the fragment + /// covered, which is the bug. Both segments are genuinely built (the commit + /// boundary cross-checks a covering declaration against the auxiliary storage, + /// so hand-doctored metadata cannot reach the prune). #[tokio::test] async fn test_prune_stale_coverage_notices_a_changed_carried_column() { - use crate::dataset::NewColumnTransform; + use crate::index::vector::VectorIndexParams; use lance_datagen::{BatchCount, RowCount, array}; let test_dir = tempfile::tempdir().unwrap(); let reader = lance_datagen::gen_batch() .col("id", array::step::()) + .col("payload", array::step::()) .col( "vector", array::rand_vec::(8.into()), ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) .await .unwrap(); let vector_field_id = dataset.schema().field("vector").unwrap().id; - // Built against the schema as it stands now, before `payload` exists. - let metadata = write_vector_segment_metadata( - &dataset, - "vector_idx", - vector_field_id, - Uuid::new_v4(), - [0_u32], - b"segment", - ) - .await; - // `payload` lands in a new data file on the same fragment, so the fragment's - // file layout changes for `payload` while `vector`'s file is untouched. + // A real covered segment carrying `payload`, and a plain control segment, + // both built against the current file layout. + let mut covered_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + covered_params.covering_columns(vec!["payload".to_string()]); + let covered = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + let plain_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + let plain = dataset + .create_index_builder(&["vector"], IndexType::Vector, &plain_params) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + + // Rewrite `payload`'s backing file IN PLACE (same fragment id, new data + // file for `payload` only); `vector`'s file is untouched. A row-moving + // update would replace the fragment and prune BOTH segments' coverage, + // which the control below would catch. + let update_schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int32, false), + arrow_schema::Field::new("payload", arrow_schema::DataType::Int32, false), + ])); + let update_batch = RecordBatch::try_new( + update_schema.clone(), + vec![ + Arc::new(arrow_array::Int32Array::from(vec![2, 4, 6])), + Arc::new(arrow_array::Int32Array::from(vec![99, 99, 99])), + ], + ) + .unwrap(); + let right: Box = + Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(update_batch)].into_iter(), + update_schema, + )); + let mut frag0 = dataset.get_fragment(0).unwrap(); + let (updated_fragment, fields_modified) = + frag0.update_columns(right, "id", "id").await.unwrap(); + let op = Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![updated_fragment], + new_fragments: vec![], + fields_modified, + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(crate::dataset::transaction::UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let read_version = dataset.manifest.version; + let mut dataset = dataset; dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("payload".into(), "id * 2".into())]), - None, - None, + .apply_commit( + Transaction::new(read_version, op, None), + &Default::default(), + &Default::default(), ) .await .unwrap(); - let payload_field_id = dataset.schema().field("payload").unwrap().id; - - let mut covered = metadata.clone(); - covered.fields = vec![vector_field_id, payload_field_id]; - covered.covering_fields = vec![payload_field_id]; let new_indices = build_index_metadata_from_segments( &dataset, @@ -9020,12 +9023,12 @@ mod tests { // Control: with nothing carried, the same segment over the same fragment // keeps its coverage -- so the assertion above is about the carried column, - // not about `add_columns` invalidating everything. + // not about the update invalidating everything. let plain_indices = build_index_metadata_from_segments( &dataset, "vector_idx", vector_field_id, - vec![segment_from_metadata(&metadata)], + vec![segment_from_metadata(&plain)], ) .await .unwrap(); @@ -9035,7 +9038,7 @@ mod tests { .as_ref() .unwrap() .is_empty(), - "a non-covering segment must keep its coverage across the same add_columns" + "a non-covering segment must keep its coverage across the same update" ); } @@ -9243,6 +9246,75 @@ mod tests { assert!(error.to_string().contains("are not among its fields")); } + /// Two segments passed to a single `build_index_metadata_from_segments` + /// call must agree on `covering_fields`: the read path (`knn.rs`, + /// `scanner.rs`) derives its declared exec output schema from a single + /// segment and assumes every sibling of the same logical index agrees. + /// Each segment individually satisfies the per-segment suffix rule + /// (`validate_covering_fields`), so only a cross-segment comparison + /// catches the disagreement. Calls `build_index_metadata_from_segments` + /// directly, not through `commit_existing_index_segments`, so the + /// assertion is evidence for this guard specifically. + #[tokio::test] + async fn test_build_index_metadata_from_segments_rejects_covering_fields_disagreement() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + let dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let id_field_id = dataset.schema().field("id").unwrap().id; + + let mut covered = write_vector_segment_metadata( + &dataset, + "vector_idx", + vector_field_id, + Uuid::new_v4(), + [0_u32], + b"seg0", + ) + .await; + covered.fields = vec![vector_field_id, id_field_id]; + covered.covering_fields = vec![id_field_id]; + + let uncovered = write_vector_segment_metadata( + &dataset, + "vector_idx", + vector_field_id, + Uuid::new_v4(), + [1_u32], + b"seg1", + ) + .await; + + let error = build_index_metadata_from_segments( + &dataset, + "vector_idx", + vector_field_id, + vec![ + segment_from_metadata(&covered), + segment_from_metadata(&uncovered), + ], + ) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("every segment of one index must declare the same columns"), + "unexpected error: {error}" + ); + } + #[tokio::test] async fn test_commit_existing_index_segments_rejects_wrong_field_provenance() { use lance_datagen::{BatchCount, RowCount, array}; diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 8bfb32aec75..adfe0f9e946 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -663,6 +663,7 @@ async fn build_fresh_vector_segment( logical_index: &LogicalVectorIndex, field_path: &str, fragment_bitmap: &RoaringBitmap, + covering_columns: &[String], progress: Arc, ) -> Result { let (reference_metadata, reference_index) = logical_index.iter().last().ok_or_else(|| { @@ -671,7 +672,12 @@ async fn build_fresh_vector_segment( logical_index.name() )) })?; - let params = fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; + let mut params = fresh_vector_segment_params(reference_metadata, reference_index.as_ref())?; + // A retrain rebuilds the storage from scratch, so re-declare the covering + // ("included") columns -- otherwise the fresh files omit the payload while the + // committed metadata still advertises `covering_fields`, and covered queries + // then expect columns the storage cannot emit. + params.covering_columns(covering_columns.to_vec()); let mut build_dataset = dataset.clone(); CreateIndexBuilder::new( &mut build_dataset, @@ -692,12 +698,17 @@ async fn scan_vector_fragments( field_path: &str, column_nullable: bool, fragments: &[Fragment], + covering_columns: &[String], ) -> Result { let mut scanner = dataset.scan(); + // Project the vector column plus any covering ("included") columns so rows from the + // newly-indexed fragments carry the same columns as the existing partitions. + let mut projection: Vec<&str> = vec![field_path]; + projection.extend(covering_columns.iter().map(String::as_str)); scanner .with_fragments(fragments.to_vec()) .with_row_id() - .project(&[field_path])?; + .project(&projection)?; if column_nullable { let column_expr = lance_datafusion::logical_expr::field_path_to_expr(field_path)?; scanner.filter_expr(column_expr.is_not_null()); @@ -788,6 +799,38 @@ pub async fn merge_indices_with_unindexed_frags<'a>( .as_ref() .map(|resolved| resolved.canonical_path.clone()) .unwrap_or(raw_field_path); + // Covering ("included") columns recorded on the index, resolved to names. + // Threaded into the merge builder and projected into the new-data scan so + // they survive optimize/merge and partition split/join. Resolution is + // TOP-LEVEL ONLY, matching every other covering resolution site (creation + // rejects dotted paths, so `covering_fields` can only hold top-level ids): + // the recursive `Schema::field_by_id` would resolve a nested id -- possible + // only in corrupt or foreign metadata -- to a leaf name and rebuild storage + // under a name the read path errors on, making commit-time and query-time + // disagree about the same index. An unresolvable id must fail loudly: the + // merged delta is committed with the original `covering_fields`, so silently + // rebuilding without the column would leave metadata advertising a column + // the storage cannot emit. + let covering_columns: Vec = old_indices[0] + .covering_fields + .iter() + .map(|id| { + dataset + .schema() + .fields + .iter() + .find(|f| f.id == *id) + .map(|f| f.name.clone()) + .ok_or_else(|| { + Error::index(format!( + "Append index: covering field id {} recorded on index '{}' does not \ + exist as a top-level field in the dataset schema; index metadata and \ + schema are inconsistent", + id, old_indices[0].name + )) + }) + }) + .collect::>()?; let first_is_vector_index = metadata_is_vector_index(dataset.as_ref(), old_indices[0]).await?; for idx in old_indices.iter().skip(1) { let is_vector_index = metadata_is_vector_index(dataset.as_ref(), idx).await?; @@ -878,6 +921,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &logical_index, &field_path, &fragment_bitmap, + &covering_columns, options.progress.clone(), ) .await?; @@ -920,6 +964,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, column.nullable, unindexed, + &covering_columns, ) .await?; let mut append_options = options.clone(); @@ -931,6 +976,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &reference_ivf_view, &append_options, + &covering_columns, ) .boxed() .await?; @@ -1012,6 +1058,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &selected_ivf_view, options, + &covering_columns, )) .await?; if indices_merged == 0 { @@ -1067,6 +1114,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, column.nullable, unindexed, + &covering_columns, ) .await?, ) @@ -1078,6 +1126,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( &field_path, &merge_ivf_view, options, + &covering_columns, ) .boxed() .await?; @@ -2954,6 +3003,69 @@ mod tests { ); } + /// An index whose `covering_fields` records a field id missing from the current + /// schema is metadata/schema-inconsistent. Optimize must fail loudly rather than + /// silently rebuild without the covered column: the merged delta would be committed + /// still advertising the column, and every covered query would then fail at read + /// time. + #[tokio::test] + async fn test_merge_indices_errors_on_unresolvable_covered_field() { + const DIM: usize = 16; + const TOTAL: usize = 256; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let vectors = generate_random_array(TOTAL * DIM); + let schema = Arc::new(Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + Field::new("id", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()), + Arc::new(UInt32Array::from_iter_values(0..TOTAL as u32)), + ], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + let index_params = VectorIndexParams::ivf_pq(2, 8, 4, MetricType::L2, 2); + dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, true) + .await + .unwrap(); + + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + let mut stale_index = indices[0].clone(); + stale_index.covering_fields = vec![9999]; + let old_indices = vec![&stale_index]; + let result = merge_indices_with_unindexed_frags( + Arc::new(dataset), + &old_indices, + &[], + &OptimizeOptions::merge(1), + ) + .await; + let err = result.unwrap_err(); + assert!( + matches!(err, Error::Index { .. }), + "unexpected error: {err}" + ); + assert!( + err.to_string().contains("9999"), + "error should name the unresolvable covered field id: {err}" + ); + } + #[tokio::test] async fn test_optimize_btree_multi_segment_optimize_default() { async fn query_id_count(dataset: &Dataset, id: &str) -> usize { @@ -3095,6 +3207,86 @@ mod tests { ); } + /// Covering resolution is top-level only (creation rejects dotted paths, so + /// `covering_fields` can only legally hold top-level ids). A NESTED field id -- + /// possible only in corrupt or foreign metadata -- must therefore fail loudly + /// here too, not resolve through the recursive `Schema::field_by_id` to a leaf + /// name and rebuild storage under a name the read path errors on: commit-time + /// and query-time must not disagree about the same index. + #[tokio::test] + async fn test_merge_indices_rejects_nested_covered_field_id() { + const DIM: usize = 16; + const TOTAL: usize = 256; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let vectors = generate_random_array(TOTAL * DIM); + let schema = Arc::new(Schema::new(vec![ + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + Field::new( + "meta", + DataType::Struct(vec![Field::new("code", DataType::UInt32, false)].into()), + false, + ), + ])); + let codes = Arc::new(UInt32Array::from_iter_values(0..TOTAL as u32)); + let meta = Arc::new(arrow_array::StructArray::from(vec![( + Arc::new(Field::new("code", DataType::UInt32, false)), + codes as arrow_array::ArrayRef, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()), + meta, + ], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + let index_params = VectorIndexParams::ivf_pq(2, 8, 4, MetricType::L2, 2); + dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, true) + .await + .unwrap(); + + // The LEAF id of meta.code -- resolvable by the recursive lookup, but not a + // top-level field. + let leaf_id = dataset + .schema() + .field("meta") + .unwrap() + .children + .first() + .unwrap() + .id; + + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + let mut nested_covered = indices[0].clone(); + nested_covered.covering_fields = vec![leaf_id]; + let old_indices = vec![&nested_covered]; + let err = merge_indices_with_unindexed_frags( + Arc::new(dataset), + &old_indices, + &[], + &OptimizeOptions::merge(1), + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("top-level"), + "a nested covered field id must be rejected as non-top-level: {err}" + ); + } + #[tokio::test] async fn test_optimize_fmindex_default_rebuilds_old_and_new_rows() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 8cd006a59a1..767c7c059fa 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -597,11 +597,41 @@ impl<'a> CreateIndexBuilder<'a> { } }; + // Resolve any covering ("included") columns to field ids for the generic + // IndexMetadata. The build path validates these names first, so a failure here is + // defence-in-depth -- but this is the site that decides what lands in the manifest, + // so it refuses rather than dropping: silently skipping a name would commit an + // index whose `fields` under-declare the payload its storage actually carries, + // leaving that column out of every invalidation rule keyed on `fields`. + let covering_fields = match self.params.as_any().downcast_ref::() { + Some(vp) => vp + .covering_columns + .iter() + .map(|name| { + self.dataset + .schema() + .field(name) + .map(|f| f.id) + .ok_or_else(|| { + Error::index(format!( + "covering column '{name}' is not present in the dataset schema" + )) + }) + }) + .collect::>>()?, + None => Vec::new(), + }; + // `covering_fields` must be the trailing entries of `fields` (see + // `IndexMetadata::validate_covering_fields`), so the keyed field comes first + // and the covering fields are appended, not tracked separately. + let mut fields = vec![field.id]; + fields.extend_from_slice(&covering_fields); + Ok(IndexMetadata { uuid: output_index_uuid, name: index_name, - fields: vec![field.id], - covering_fields: vec![], + fields, + covering_fields, dataset_version: self.dataset.manifest.version, fragment_bitmap: if train { match &self.fragments { diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index d9e749d44ad..8d57d6893fb 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -276,6 +276,10 @@ pub struct VectorIndexParams { /// Keys use reverse-DNS namespacing (e.g., "lance.ivf.max_iters"). /// Populated by the build path and merged into VectorIndexDetails at creation time. pub runtime_hints: HashMap, + + /// Columns to co-locate ("include") in the index storage alongside the row id + /// and the quantization code. Empty by default. + pub covering_columns: Vec, } impl VectorIndexParams { @@ -289,6 +293,12 @@ impl VectorIndexParams { self } + /// Set the columns to co-locate ("include") in the index storage. + pub fn covering_columns(&mut self, columns: Vec) -> &mut Self { + self.covering_columns = columns; + self + } + pub fn ivf_flat(num_partitions: usize, metric_type: MetricType) -> Self { let ivf_params = IvfBuildParams::new(num_partitions); let stages = vec![StageParams::Ivf(ivf_params)]; @@ -298,6 +308,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -309,6 +320,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -344,6 +356,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -371,6 +384,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -387,6 +401,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -402,6 +417,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -417,6 +433,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -432,6 +449,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -454,6 +472,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -476,6 +495,7 @@ impl VectorIndexParams { version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: HashMap::new(), + covering_columns: Vec::new(), } } @@ -612,6 +632,18 @@ pub(crate) async fn build_distributed_vector_index( fragment_ids: &[u32], progress: Arc, ) -> Result<(Uuid, Vec)> { + // The distributed shard builders and the distributed merger have no covering-column + // plumbing: each shard would write storage without the payload while the committed + // metadata still advertises `covering_fields`, desyncing the two. Reject covering on + // this path up front rather than publishing an inconsistent index. + if !params.covering_columns.is_empty() { + return Err(Error::invalid_input( + "covering_columns (covering columns) are not supported for distributed vector index \ + builds (precomputed IVF centroids with a fragment subset). Build the covered index \ + without a fragment restriction / precomputed IVF." + .to_string(), + )); + } let (element_type, index_type, ivf_params, shuffler) = prepare_vector_segment_build( dataset, column, @@ -1025,6 +1057,97 @@ async fn build_vector_index_impl( .await?; let stages = ¶ms.stages; + // Covering ("included") columns require the V3 index file format: only the V3 + // covering-aware storages preserve the extra columns row-aligned and report them via + // `covering_field_indices`. Every vector index type supports covering at V3, so reject + // only legacy/non-V3 builds up front -- otherwise the build silently ignores the option + // while `covering_fields` still lands in the index metadata, and the search exec then + // declares a covered schema the storage cannot emit. + if !params.covering_columns.is_empty() && !matches!(params.version, IndexFileVersion::V3) { + return Err(Error::invalid_input(format!( + "covering_columns (covering columns) requires index file version V3, but got {:?}", + params.version + ))); + } + + // A covering column is stored inline in the index, row-aligned with the code, and + // projected by name from the source batch. Reject columns that cannot satisfy that + // contract up front (otherwise the build fails deep in projection, or -- worse -- + // stores unusable data): nested/dotted paths (the covering gather projects only + // top-level names) and blob columns (which store out-of-line descriptors, not inline + // data). + // Names that collide with a build pipeline's own internal columns. Covering one would + // advertise a column in `covering_fields` that the storage's `covering_field_indices` + // excludes (so a covered query declares a column storage never emits), or fail deep in + // the quantizer transform. This check runs before the `match index_type` below, so it + // must list the *union* of every pipeline's internal names -- the row id, distance, + // partition id, each quantizer's code column, RaBitQ's extended-code and per-row factor + // columns, and the IVF partition transform's transient `__centroid_dist`. + const RESERVED_STORAGE_COLUMNS: &[&str] = &[ + lance_core::ROW_ID, + lance_index::vector::DIST_COL, + lance_index::vector::PART_ID_COLUMN, + lance_index::vector::CENTROID_DIST_COLUMN, + lance_index::vector::PQ_CODE_COLUMN, + lance_index::vector::SQ_CODE_COLUMN, + lance_index::vector::flat::storage::FLAT_COLUMN, + lance_index::vector::bq::storage::RABIT_CODE_COLUMN, + lance_index::vector::bq::storage::RABIT_EX_CODE_COLUMN, + lance_index::vector::bq::storage::RABIT_BLOCKED_EX_CODE_COLUMN, + lance_index::vector::bq::transform::ADD_FACTORS_COLUMN, + lance_index::vector::bq::transform::SCALE_FACTORS_COLUMN, + lance_index::vector::bq::transform::ERROR_FACTORS_COLUMN, + lance_index::vector::bq::transform::EX_ADD_FACTORS_COLUMN, + lance_index::vector::bq::transform::EX_SCALE_FACTORS_COLUMN, + ]; + let mut seen_include = std::collections::HashSet::with_capacity(params.covering_columns.len()); + for name in ¶ms.covering_columns { + if name.contains('.') { + return Err(Error::invalid_input(format!( + "covering_columns: nested/dotted covering column '{name}' is not supported; only \ + top-level columns can be covered" + ))); + } + if name == column { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is the indexed vector column itself; it \ + is stored as the quantization code, not a covered payload" + ))); + } + if RESERVED_STORAGE_COLUMNS.contains(&name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' collides with a reserved index storage \ + column name" + ))); + } + if !seen_include.insert(name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: duplicate covering column '{name}'" + ))); + } + let field = dataset.schema().field(name).ok_or_else(|| { + Error::invalid_input(format!( + "covering_columns: covering column '{name}' does not exist in the dataset schema" + )) + })?; + if field.is_blob() { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is a blob column; blob columns cannot be \ + covered by a vector index" + ))); + } + } + + // Covering is implemented for IVF_PQ in this change; the other IVF vector types land in a + // follow-up. Reject covering on a non-PQ index up front rather than recording + // `covering_fields` the storage cannot serve. + if !params.covering_columns.is_empty() && !matches!(index_type, IndexType::IvfPq) { + return Err(Error::invalid_input( + "covering_columns (covering columns) are currently only supported for IVF_PQ indexes" + .to_string(), + )); + } + match index_type { IndexType::IvfFlat => match element_type { DataType::Float16 | DataType::Float32 | DataType::Float64 => { @@ -1039,6 +1162,7 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1057,6 +1181,7 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1112,6 +1237,7 @@ async fn build_vector_index_impl( )?; let summary = builder + .with_covering_columns(params.covering_columns.clone()) .with_transpose(!params.skip_transpose) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) @@ -1140,6 +1266,7 @@ async fn build_vector_index_impl( (), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1167,6 +1294,7 @@ async fn build_vector_index_impl( )?; let summary = builder + .with_covering_columns(params.covering_columns.clone()) .with_transpose(!params.skip_transpose) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) @@ -1194,6 +1322,7 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1212,6 +1341,7 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1244,6 +1374,7 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1274,6 +1405,7 @@ async fn build_vector_index_impl( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_optional_fragment_filter(fragment_ids) .with_progress(progress.clone()) .build() @@ -1355,6 +1487,11 @@ pub(crate) async fn build_vector_index_incremental( // Determine the index type and build incrementally let (sub_index_type, quantization_type) = existing_index.sub_index_type(); + // Every vector index type supports covering ("included") columns, so the + // incremental/copy path needs no per-type gate here -- each branch below threads + // `covering_columns` into its builder, and the source index was already built with a + // covering-aware storage. + match (sub_index_type, quantization_type) { // IVF_FLAT (SubIndexType::Flat, QuantizationType::Flat) => { @@ -1370,6 +1507,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1388,6 +1526,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1409,6 +1548,7 @@ pub(crate) async fn build_vector_index_incremental( .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) .with_transpose(!params.skip_transpose) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1428,6 +1568,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1448,6 +1589,7 @@ pub(crate) async fn build_vector_index_incremental( let summary = builder .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_transpose(!params.skip_transpose) .with_progress(progress.clone()) .build() @@ -1477,6 +1619,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1495,6 +1638,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1513,6 +1657,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1531,6 +1676,7 @@ pub(crate) async fn build_vector_index_incremental( )? .with_ivf(ivf_model) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(params.covering_columns.clone()) .with_progress(progress.clone()) .build() .await?; @@ -1864,7 +2010,7 @@ pub async fn initialize_vector_index( let (sub_index_type, quantization_type) = source_vector_index.sub_index_type(); let ivf_params = derive_ivf_params(ivf_model); - let params = match (sub_index_type, quantization_type) { + let mut params = match (sub_index_type, quantization_type) { (SubIndexType::Flat, QuantizationType::Flat) | (SubIndexType::Flat, QuantizationType::FlatBin) => { VectorIndexParams::with_ivf_flat_params(metric_type, ivf_params) @@ -1919,6 +2065,27 @@ pub async fn initialize_vector_index( } }; + // Carry the source index's covering ("included") columns so the rebuilt + // target storage keeps them. covering_fields are field ids in the source; + // resolve to names (the same columns must exist in the target dataset). + let covering_columns: Vec = source_index + .covering_fields + .iter() + .map(|id| { + source_dataset + .schema() + .field_by_id(*id) + .map(|f| f.name.clone()) + .ok_or_else(|| { + Error::index(format!( + "covering field id {id} (from index '{}') not found in source dataset schema", + source_index.name + )) + }) + }) + .collect::>>()?; + params.covering_columns(covering_columns); + let new_uuid = Uuid::new_v4(); let frag_reuse_index = target_dataset .open_frag_reuse_index(&NoOpMetricsCollector) @@ -1944,11 +2111,35 @@ pub async fn initialize_vector_index( let fragment_bitmap = Some(target_dataset.fragment_bitmap.as_ref().clone()); + // Field ids are per-dataset: re-resolve the covering column names against + // the TARGET schema (like `fields` below), never copy the source's ids. + let covering_fields = params + .covering_columns + .iter() + .map(|name| { + target_dataset + .schema() + .field(name) + .map(|f| f.id) + .ok_or_else(|| { + Error::index(format!( + "covering column '{}' not found in target dataset schema", + name + )) + }) + }) + .collect::>>()?; + // `covering_fields` must be the trailing entries of `fields` (see + // `IndexMetadata::validate_covering_fields`), so the keyed field comes first + // and the covering fields are appended, not tracked separately. + let mut fields = vec![field.id]; + fields.extend_from_slice(&covering_fields); + let new_idx = IndexMetadata { uuid: new_uuid, name: source_index.name.clone(), - fields: vec![field.id], - covering_fields: vec![], + fields, + covering_fields, dataset_version: target_dataset.manifest.version, fragment_bitmap, index_details: source_index.index_details.clone(), @@ -2525,6 +2716,211 @@ mod tests { assert_eq!(results.num_rows(), 10, "Should return 10 nearest neighbors"); } + /// `covering_columns` requires the V3 index file format. A legacy/non-V3 build must be + /// rejected at create time: the build would silently ignore the option while + /// `create_index` still records `covering_fields` in the manifest, so the exec would + /// declare a covered schema the storage cannot emit and every query on the index would + /// fail. + /// + /// Covering is implemented for IVF_PQ only in this change (see the covering validation + /// in `build_vector_index_impl`); every other IVF vector type must reject it too, or a + /// covered index of that type commits `covering_fields` its storage cannot serve and + /// every covered query on it fails at read time instead of at creation time. These + /// non-PQ cases are exactly the ones a would-be `test_initialize_vector_index_preserves_covering` + /// parametrization would need once covering extends beyond IVF_PQ -- until then, they + /// belong here, asserting rejection, not there asserting success. + #[tokio::test] + async fn test_covering_columns_rejected_for_unsupported_index_types() { + let test_dir = TempStrDir::default(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.as_str(), None) + .await + .unwrap(); + + const ONLY_IVF_PQ: &str = "only supported for IVF_PQ"; + for (label, mut params, expected_fragment) in [ + ( + "IVF_PQ legacy", + VectorIndexParams::with_ivf_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + PQBuildParams::new(8, 8), + ) + .version(IndexFileVersion::Legacy) + .clone(), + "covering_columns", + ), + ( + "IVF_SQ", + VectorIndexParams::with_ivf_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + SQBuildParams::default(), + ), + ONLY_IVF_PQ, + ), + ( + "IVF_HNSW_PQ", + VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(8, 8), + ), + ONLY_IVF_PQ, + ), + ( + "IVF_HNSW_SQ", + VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default(), + ), + ONLY_IVF_PQ, + ), + ( + "IVF_RQ", + VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ), + ONLY_IVF_PQ, + ), + ( + "IVF_FLAT", + VectorIndexParams::ivf_flat(4, MetricType::L2), + ONLY_IVF_PQ, + ), + ( + "IVF_HNSW_FLAT", + VectorIndexParams::ivf_hnsw( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + ), + ONLY_IVF_PQ, + ), + ] { + params.covering_columns(vec!["id".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err(&format!("covering_columns on {label} must be rejected")); + assert!( + err.to_string().contains(expected_fragment), + "{label}: error should contain {expected_fragment:?}, got: {err}" + ); + // Fail fast: no index metadata must have been committed. + assert!( + dataset.load_indices().await.unwrap().is_empty(), + "{label}: no index should have been created" + ); + } + } + + /// Copying/importing a covered index into another dataset (initialize) must + /// rebuild the target storage WITH the covering columns, not just copy the + /// `covering_fields` metadata -- otherwise the target advertises covering + /// columns its storage can't emit and covered queries fail. + /// + /// IVF_PQ only: covering is implemented only for IVF_PQ in this change (see + /// `build_vector_index_impl`'s covering validation), so this only parametrizes + /// over PQ. Other IVF vector types land coverage in a follow-up that extends + /// `covering_columns` support to them. + #[rstest::rstest] + #[case::pq(VectorIndexParams::ivf_pq(10, 8, 16, MetricType::L2, 50))] + #[tokio::test] + async fn test_initialize_vector_index_preserves_covering( + #[case] mut params: VectorIndexParams, + ) { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + 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)); + let mut source_dataset = Dataset::write(source_reader, &source_uri, None) + .await + .unwrap(); + + // Covered index on source (covers `id`). + params.covering_columns(vec!["id".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + + // The target has an extra leading column, so its field ids are shifted + // relative to the source's: `covering_fields` must be re-resolved + // against the TARGET schema (by name), not copied from the source. + let target_reader = lance_datagen::gen_batch() + .col("pad", array::step::()) + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + + initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .unwrap(); + + // Metadata carried over, resolved to the target's field ids. + let target_index = target_dataset.load_indices().await.unwrap()[0].clone(); + let id_field = target_dataset.schema().field("id").unwrap().id; + assert_ne!( + id_field, source_index.covering_fields[0], + "test setup must give 'id' different field ids in source and target" + ); + assert_eq!(target_index.covering_fields, vec![id_field]); + + // End-to-end: a covered projection on the copied index skips the take and + // returns `id` -- which only works if the rebuilt storage carries it. + let query = arrow_array::Float32Array::from(vec![0.5f32; 32]); + let mut scan = target_dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection should skip the take on the copied index; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.column_by_name("id").is_some(), + "copied index should emit the covering column 'id'" + ); + } + #[tokio::test] async fn test_initialize_vector_index_ivf_flat() { let test_dir = TempStrDir::default(); @@ -2776,6 +3172,47 @@ mod tests { ); } + /// The distributed build + merge path has no covering-column plumbing, so it must + /// reject `covering_columns` up front rather than publish an index whose metadata + /// advertises columns the shard files do not contain. + #[tokio::test] + async fn test_build_distributed_rejects_covering_columns() { + let test_dir = TempStrDir::default(); + let uri = format!("{}/ds", test_dir.as_str()); + + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(128), BatchCount::from(1)); + let dataset = Dataset::write(reader, &uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(4, MetricType::L2); + params.covering_columns(vec!["id".to_string()]); + + let result = build_distributed_vector_index( + &dataset, + "vector", + "vector_dist", + Uuid::new_v4(), + ¶ms, + None, + &[0], + noop_progress(), + ) + .await; + + assert!( + matches!(&result, Err(Error::InvalidInput { .. })), + "distributed build must reject covering_columns, got {:?}", + result + ); + let msg = result.unwrap_err().to_string(); + assert!( + msg.contains("covering_columns") && msg.contains("distributed"), + "error should mention covering_columns and the distributed path; got: {msg}" + ); + } + #[tokio::test] async fn test_build_distributed_empty_fragment_ids() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 3a65d4776d1..3121574732d 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -78,7 +78,7 @@ use lance_table::format::IndexFile; use log::info; use object_store::path::Path; use prost::Message; -use roaring::RoaringBitmap; +use roaring::{RoaringBitmap, RoaringTreemap}; use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore}; use tracing::{Level, instrument, span}; @@ -94,7 +94,7 @@ use crate::index::vector::utils::infer_vector_dim; use super::v2::IVFIndex; use super::{ ivf::load_precomputed_partitions_if_available, - utils::{self, get_vector_type}, + utils::{self, gather_covering_columns_by_row_id, get_vector_type}, }; // the number of partitions to evaluate for reassigning @@ -122,6 +122,22 @@ impl Default for FreshPartitionBuildLimits { } } +/// Append the covering columns from `covering` (`[_rowid, ]`) +/// onto `batch` (which must contain a `_rowid` column), gathered by row id. +/// Used to re-attach covering columns to rows that were reordered/regrouped by +/// the partition-join reassignment. +fn append_covering_by_row_id( + mut batch: RecordBatch, + covering: &RecordBatch, +) -> Result { + let target_ids = batch[ROW_ID].as_primitive::(); + let included = gather_covering_columns_by_row_id(covering, target_ids)?; + for (field, col) in included { + batch = batch.try_with_column(field.as_ref().clone(), col)?; + } + Ok(batch) +} + /// Build a new centroid array that incorporates the results of partition splits. /// /// For each `(part_idx, centroid1, centroid2)` in `splits`: @@ -247,6 +263,9 @@ pub struct IvfIndexBuilder { ivf_params: Option, quantizer_params: Option, sub_index_params: Option, + /// Columns to co-locate ("include") in the index storage alongside the row + /// id and quantization code, so covered queries can avoid a take. + covering_columns: Vec, _temp_dir: TempStdDir, // store this for keeping the temp dir alive and clean up after build temp_dir: Path, @@ -376,6 +395,7 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + covering_columns: Vec::new(), format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -418,6 +438,34 @@ impl IvfIndexBuilder .downcast_ref::>() .ok_or(Error::invalid_input("existing index is not IVF index"))?; + // Carry the covering ("included") columns so the remapped storage keeps + // them -- otherwise merge_partitions would write a schema without them + // while the metadata still advertises them. + // + // This set comes from *storage*, so in principle it could name a column the + // schema no longer has. Refuse rather than silently dropping it: `remap_index` + // copies `covering_fields` onto the replacement metadata verbatim + // (`dataset/optimize/remapping.rs`), so dropping the column here would commit an + // index whose declaration its storage cannot satisfy, and every later query would + // fail in `effective_covering` ("declares covering field id N ...") rather than + // falling back. A loud failure here beats a silent one at read time. + // + // Unreachable on the commit path today: covering fields live in `fields`, and + // `retain_relevant_indices` drops any index with a field missing from the schema, + // so dropping a covered column deletes the whole index instead of stranding it. + // Defence-in-depth against a manifest this build did not write. + let covering_columns = ivf_index.covering_column_names()?; + if let Some(missing) = covering_columns + .iter() + .find(|name| dataset.schema().field(name).is_none()) + { + return Err(Error::index(format!( + "cannot remap index: its storage carries covering column '{missing}', which \ + is not present in the current dataset schema; index metadata and schema \ + are inconsistent (covering columns: {covering_columns:?})" + ))); + } + let temp_dir = TempStdDir::default(); let temp_dir_path = Path::from_filesystem_path(&temp_dir)?; let format_version = dataset_format_version(&dataset); @@ -443,6 +491,7 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + covering_columns, format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -579,6 +628,13 @@ impl IvfIndexBuilder self } + /// Set columns to co-locate ("include") in the index storage so covered + /// queries can avoid a take from the base table. + pub fn with_covering_columns(&mut self, columns: Vec) -> &mut Self { + self.covering_columns = columns; + self + } + /// Set progress callback for index building pub fn with_progress(&mut self, progress: Arc) -> &mut Self { self.progress = progress; @@ -746,6 +802,16 @@ impl IvfIndexBuilder .and_then(|p| p.precomputed_shuffle_buffers.as_ref()) { Some((uri, _)) => { + // Precomputed shuffle buffers were materialized without the covering + // columns, so a covered build reading from them would silently omit those + // columns from partition storage. Reject the combination up front. + if !self.covering_columns.is_empty() { + return Err(Error::invalid_input( + "covering_columns (covering columns) are not supported with precomputed \ + shuffle buffers: the precomputed buffers do not carry the covering columns" + .to_string(), + )); + } let uri = to_local_path(uri); // the uri points to data directory, // so need to trim the "data" suffix for reading the dataset @@ -756,10 +822,15 @@ impl IvfIndexBuilder } _ => { log::info!("shuffle column {} over dataset", self.column); + // Read the vector column plus any included ("covering") columns + // so they flow through the shuffle into the partition storage. + let mut projection: Vec<&str> = Vec::with_capacity(1 + self.covering_columns.len()); + projection.push(self.column.as_str()); + projection.extend(self.covering_columns.iter().map(String::as_str)); let mut builder = dataset.scan(); builder .batch_readahead(get_num_compute_intensive_cpus()) - .project(&[self.column.as_str()])? + .project(&projection)? .with_row_id(); // Apply fragment filter for distributed indexing @@ -1072,6 +1143,9 @@ impl IvfIndexBuilder FreshPartitionBuildLimits::default(), ); } + // Covering builds may re-scan a pruned fragment that still lives in old + // storage; drop the stale duplicate row ids during the partition read. + let dedup_existing = !self.covering_columns.is_empty(); let partition_adjustment = Arc::new(partition_adjustment); let build_iter = assign_batches @@ -1086,6 +1160,7 @@ impl IvfIndexBuilder let sub_index_params = sub_index_params.clone(); let column = column.clone(); let frag_reuse_index = frag_reuse_index.clone(); + let dedup_existing = dedup_existing; let partition_adjustment = partition_adjustment.clone(); async move { let (is_affected, split_reader) = match partition_adjustment.as_ref() { @@ -1116,6 +1191,7 @@ impl IvfIndexBuilder partition, &[], Some(split_reader.as_ref().unwrap().as_ref()), + dedup_existing, ) .await? } else { @@ -1123,6 +1199,7 @@ impl IvfIndexBuilder partition, indices.as_ref(), Some(reader.as_ref()), + dedup_existing, ) .await? }; @@ -1130,9 +1207,13 @@ impl IvfIndexBuilder // For unaffected partitions during a split, vectors from // affected partitions may have been reassigned here. if !is_affected && let Some(sr) = split_reader.as_ref() { - let (extra, extra_loss) = - Self::take_partition_batches(partition, &[], Some(sr.as_ref())) - .await?; + let (extra, extra_loss) = Self::take_partition_batches( + partition, + &[], + Some(sr.as_ref()), + dedup_existing, + ) + .await?; batches.extend(extra); loss += extra_loss; } @@ -1424,6 +1505,7 @@ impl IvfIndexBuilder part_id: usize, existing_indices: &[ExistingIndex], reader: Option<&dyn ShuffleReader>, + dedup_by_row_id: bool, ) -> Result<(Vec, f64)> { let mut batches = Vec::new(); for source in existing_indices.iter() { @@ -1500,6 +1582,11 @@ impl IvfIndexBuilder batches.extend(part_batches); } + // Everything pushed so far came from existing-segment storage; anything + // pushed below comes from the fresh shuffle reader. The dedup at the end + // only removes stale existing rows whose row id reappears in the fresh data. + let existing_end = batches.len(); + let mut loss = 0.0; // Skip if this partition doesn't exist in the reader // This can happen after a split creates a new partition @@ -1524,9 +1611,80 @@ impl IvfIndexBuilder } } + // When rebuilding a covering index, a fragment pruned from a segment's + // coverage (e.g. an in-place rewrite of a covering column) is re-scanned as + // fresh unindexed data *and* is still present, stale, in an old segment's + // storage -- the same row id then appears in both. Drop the stale existing + // copy in favour of the fresh one. + // + // The dedup is scoped to the existing-storage <-> fresh-reader boundary: it + // removes only an *existing-storage* row whose row id reappears in the fresh + // reader data, and never dedups within a single source. That preserves + // legitimate repeated row ids -- a multivector column stores one entry per + // sub-vector, all sharing the source row's id -- which a global per-row-id + // dedup would otherwise silently collapse. + if dedup_by_row_id && existing_end > 0 { + // A roaring treemap rather than a `HashSet`: this holds every fresh row + // id in the partition, HNSW partitions target `1 << 20` rows, and + // `build_partitions` runs several partitions concurrently. + let mut fresh = RoaringTreemap::new(); + for batch in &batches[existing_end..] { + for rid in batch[ROW_ID].as_primitive::().iter().flatten() { + fresh.insert(rid); + } + } + if !fresh.is_empty() { + for batch in batches[..existing_end].iter_mut() { + if batch.num_rows() == 0 { + continue; + } + let row_ids = batch[ROW_ID].as_primitive::(); + let keep: Vec = row_ids + .iter() + .map(|row_id| match row_id { + Some(rid) => !fresh.contains(rid), + None => true, + }) + .collect(); + *batch = arrow::compute::filter_record_batch(batch, &BooleanArray::from(keep))?; + } + } + } + Ok((batches, loss)) } + /// Resolve the covering ("included") columns to their arrow fields from the dataset + /// schema, in declaration order. Empty when no covering columns are configured. Both + /// the code-storage schema and the empty-flat fallback schema append these so covered + /// storage is written consistently regardless of whether any partition held data. + fn covering_arrow_fields(&self) -> Result> { + if self.covering_columns.is_empty() { + return Ok(Vec::new()); + } + let Some(ds) = self.dataset.as_ref() else { + // `covering_columns` is non-empty here, so silently returning no + // fields would let the storage writer (which resolves covering + // columns by name) write storage without them while the committed + // metadata still advertises `covering_fields` -- the exact + // metadata/storage desync every other covering guard prevents. + return Err(Error::invalid_input( + "dataset not set before resolving covering columns".to_string(), + )); + }; + let ds_schema = arrow_schema::Schema::from(ds.schema()); + self.covering_columns + .iter() + .map(|name| { + ds_schema.field_with_name(name).cloned().map_err(|e| { + Error::invalid_input(format!( + "include column '{name}' not found in dataset schema: {e}" + )) + }) + }) + .collect() + } + #[instrument(name = "merge_partitions", level = "debug", skip_all)] async fn merge_partitions( &mut self, @@ -1551,11 +1709,17 @@ impl IvfIndexBuilder let index_path = self.index_dir.clone().join(INDEX_FILE_NAME); let writer_options = FileWriterOptions::default(); + // Covering columns are appended to whichever storage schema this build writes + // (code storage below, or the empty-flat fallback further down). + let covering_fields = self.covering_arrow_fields()?; let mut storage_writer = if is_flat { None } else { let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; fields.extend(quantizer.extra_fields()); + // Append any included ("covering") columns so they are persisted in + // the auxiliary storage file alongside the row id and code. + fields.extend(covering_fields.iter().cloned()); let storage_schema: Schema = (&arrow_schema::Schema::new(fields)).try_into()?; Some(file_versions::create_writer( self.format_version, @@ -1701,7 +1865,7 @@ impl IvfIndexBuilder "flat storage writer could not infer schema from empty partitions without IVF centroids", )); }; - let flat_schema = arrow_schema::Schema::new(vec![ + let mut flat_fields = vec![ ROW_ID_FIELD.as_ref().clone(), arrow_schema::Field::new( lance_index::vector::flat::storage::FLAT_COLUMN, @@ -1715,7 +1879,11 @@ impl IvfIndexBuilder ), true, ), - ]); + ]; + // An all-empty covered flat index must still declare its covering columns, + // or the written storage schema disagrees with `covering_fields`. + flat_fields.extend(covering_fields.iter().cloned()); + let flat_schema = arrow_schema::Schema::new(flat_fields); let storage_schema: Schema = (&flat_schema).try_into()?; storage_writer = Some(file_versions::create_writer( self.format_version, @@ -1805,14 +1973,17 @@ impl IvfIndexBuilder // take raw vectors from the dataset // - // returns batches of schema | row_id | vector | + // returns batches of schema | row_id | vector | | async fn take_vectors( dataset: &Dataset, column: &str, store: &ObjectStore, row_ids: &[u64], + covering_columns: &[String], ) -> Result> { - let projection = Arc::new(dataset.schema().project(&[column])?); + let mut proj_columns: Vec<&str> = vec![column]; + proj_columns.extend(covering_columns.iter().map(String::as_str)); + let projection = Arc::new(dataset.schema().project(&proj_columns)?); // arrow uses i32 for index, so we chunk the row ids to avoid large batch causing overflow let mut batches = Vec::new(); let row_ids = dataset.filter_deleted_ids(row_ids).await?; @@ -1836,11 +2007,13 @@ impl IvfIndexBuilder Ok(batches) } - // helper to load row ids and vectors for a partition + // helper to load row ids and vectors for a partition, plus a covering batch + // `[_rowid, ]` (None when the index has no covering columns) + // so the columns can be re-attached to reassigned rows after quantization. async fn load_partition_raw_vectors( &self, part_idx: usize, - ) -> Result> { + ) -> Result)>> { let Some(dataset) = self.dataset.as_ref() else { return Err(Error::invalid_input( "dataset not set before split partition", @@ -1854,11 +2027,42 @@ impl IvfIndexBuilder // dedup is needed if it's multivector row_ids.dedup(); - let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids).await?; + let batches = Self::take_vectors( + dataset, + &self.column, + &self.store, + &row_ids, + &self.covering_columns, + ) + .await?; if batches.is_empty() { return Ok(None); } let batch = arrow::compute::concat_batches(&batches[0].schema(), batches.iter())?; + // Extract the covering columns BEFORE flattening. `Flatten` replicates every + // non-partition column across a multivector row's expansion, so taking them + // afterwards would yield one copy per sub-vector; the pre-flatten batch has one + // row per row id, which is what the later gather-by-row-id expects. + let covering = if self.covering_columns.is_empty() { + None + } else { + let mut indices = vec![batch.schema().index_of(ROW_ID)?]; + for name in &self.covering_columns { + indices.push(batch.schema().index_of(name)?); + } + Some(batch.project(&indices)?) + }; + // Narrow to `[_rowid, vector]` first: `Flatten` replicates every remaining column + // across the multivector expansion, so leaving the covering columns in would + // materialize a second copy of the payload that nothing reads -- the gather below + // uses the pre-flatten `covering` batch. Skipped when the vector column is not a + // plain top-level column (precomputed buffers), where `Flatten` is a no-op anyway. + let batch = match (&covering, batch.schema().index_of(&self.column)) { + (Some(_), Ok(vector_idx)) => { + batch.project(&[batch.schema().index_of(ROW_ID)?, vector_idx])? + } + _ => batch, + }; // for multivector, we need to flatten the vectors let batch = Flatten::new(&self.column).transform(&batch)?; // need to retrieve the row ids from the batch because some rows may have been deleted @@ -1872,7 +2076,7 @@ impl IvfIndexBuilder )))? .as_fixed_size_list() .clone(); - Ok(Some((row_ids, vectors))) + Ok(Some((row_ids, vectors, covering))) } // check whether need to split or join partition @@ -2064,8 +2268,11 @@ impl IvfIndexBuilder all_row_ids.sort(); all_row_ids.dedup(); - // Stream raw vectors in chunks - let projection = Arc::new(dataset.schema().project(&[self.column.as_str()])?); + // Stream raw vectors plus any covering ("included") columns so a + // partition split preserves them in the re-quantized storage. + let mut split_projection: Vec<&str> = vec![self.column.as_str()]; + split_projection.extend(self.covering_columns.iter().map(String::as_str)); + let projection = Arc::new(dataset.schema().project(&split_projection)?); let row_ids = dataset.filter_deleted_ids(&all_row_ids).await?; let block_size = self.store.block_size(); let column = self.column.clone(); @@ -2174,7 +2381,8 @@ impl IvfIndexBuilder row_ids = (0..sample_size).map(|i| row_ids[i * stride]).collect(); } - let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids).await?; + // Centroid sampling only needs the vectors, not covering columns. + let batches = Self::take_vectors(dataset, &self.column, &self.store, &row_ids, &[]).await?; if batches.is_empty() { return Ok(None); } @@ -2251,8 +2459,9 @@ impl IvfIndexBuilder let new_centroids = FixedSizeListArray::try_new_from_values(new_centroids, centroids.value_length())?; - // take the raw vectors from dataset - let Some((row_ids, vectors)) = self.load_partition_raw_vectors(part_idx).await? else { + // take the raw vectors from dataset (plus any covering columns) + let Some((row_ids, vectors, covering)) = self.load_partition_raw_vectors(part_idx).await? + else { return Ok(AssignResult { assign_batches: vec![None; ivf.num_partitions() - 1], new_centroids, @@ -2266,6 +2475,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2276,6 +2486,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2286,6 +2497,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2296,6 +2508,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2313,6 +2526,9 @@ impl IvfIndexBuilder ivf: &IvfModel, row_ids: &UInt64Array, vectors: &FixedSizeListArray, + // `[_rowid, ]` for the reassigned rows, or None when the + // index has no covering columns. Re-attached to each assign batch by row id. + covering: Option<&RecordBatch>, new_centroids: FixedSizeListArray, ) -> Result where @@ -2356,6 +2572,21 @@ impl IvfIndexBuilder } let assign_batches = self.build_assign_batch::(&new_centroids, &assign_ops)?; + // Re-attach the covering columns to each reassigned batch, gathered by + // row id (in-memory; the rows were reordered across target partitions). + let assign_batches = match covering { + None => assign_batches, + Some(covering) => assign_batches + .into_iter() + .map(|entry| match entry { + Some((batch, deleted)) => { + Ok(Some((append_covering_by_row_id(batch, covering)?, deleted))) + } + None => Ok(None), + }) + .collect::>>()?, + }; + Ok(AssignResult { assign_batches, new_centroids, @@ -3117,6 +3348,7 @@ mod tests { 0, &[], Some(&reader), + false, ) .await .unwrap(); diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 2ac0760caf6..a7745245e70 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -649,6 +649,9 @@ pub(crate) async fn optimize_vector_indices( vector_column: &str, logical_index: &LogicalIvfView<'_>, options: &OptimizeOptions, + // Covering ("included") column names to preserve through the merge. Empty + // for ordinary indexes. Only the v2 path supports covering columns. + covering_columns: &[String], ) -> Result<(Uuid, usize, Vec)> { let existing_indices = logical_index.indices().cloned().collect::>(); // Sanity check the indices @@ -663,8 +666,15 @@ pub(crate) async fn optimize_vector_indices( // fallback to v2 IVFIndex if it's not v1 IVFIndex if !existing_indices[0].as_any().is::() { let sources = existing_index_sources(&dataset, logical_index); - return optimize_vector_indices_v2(&dataset, unindexed, vector_column, &sources, options) - .await; + return optimize_vector_indices_v2( + &dataset, + unindexed, + vector_column, + &sources, + options, + covering_columns, + ) + .await; } let new_uuid = Uuid::new_v4(); @@ -735,6 +745,9 @@ pub(crate) async fn optimize_vector_indices_v2( vector_column: &str, existing_indices: &[ExistingIndex], options: &OptimizeOptions, + // Covering ("included") column names to preserve through the merge/split/join. + // Empty for ordinary indexes. + covering_columns: &[String], ) -> Result<(Uuid, usize, Vec)> { // Sanity check the indices if existing_indices.is_empty() { @@ -761,6 +774,8 @@ pub(crate) async fn optimize_vector_indices_v2( let shuffler = create_ivf_shuffler(temp_dir_path, num_partitions, format_version, None); let (_, element_type) = get_vector_type(dataset.schema(), vector_column)?; + // Every vector index type supports covering ("included") columns, so each branch below + // threads `covering_columns` into its builder regardless of quantizer. let summary = match index_type { // IVF_FLAT (SubIndexType::Flat, QuantizationType::Flat) => { @@ -777,6 +792,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(covering_columns.to_vec()) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -795,6 +811,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(covering_columns.to_vec()) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -818,6 +835,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -838,6 +856,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -858,6 +877,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -877,6 +897,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -896,6 +917,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(covering_columns.to_vec()) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -914,6 +936,7 @@ pub(crate) async fn optimize_vector_indices_v2( )? .with_ivf(ivf_model.clone()) .with_quantizer(quantizer.try_into()?) + .with_covering_columns(covering_columns.to_vec()) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) .shuffle_data_input(unindexed) @@ -937,6 +960,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -957,6 +981,7 @@ pub(crate) async fn optimize_vector_indices_v2( .with_quantizer(quantizer.try_into()?) .with_existing_index_sources(existing_indices.clone()) .with_progress(options.progress.clone()) + .with_covering_columns(covering_columns.to_vec()) .shuffle_data_input(unindexed) .build() .await? @@ -4967,6 +4992,7 @@ mod tests { "vector", &sources, &OptimizeOptions::new(), + &[], ) .await .unwrap(); @@ -4982,6 +5008,7 @@ mod tests { "vector", &sources, &OptimizeOptions::merge(1), + &[], ) .await .unwrap(); diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index e369ad53888..3bd0f34fad5 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -8,18 +8,20 @@ use std::marker::PhantomData; use std::{ any::Any, borrow::Cow, - collections::{BinaryHeap, HashMap}, + collections::{BinaryHeap, HashMap, HashSet}, sync::{ Arc, LazyLock, Mutex, OnceLock, atomic::{AtomicBool, Ordering}, }, }; -use crate::index::vector::{IndexFileVersion, builder::index_type_string}; +use crate::index::vector::{ + IndexFileVersion, builder::index_type_string, utils::gather_covering_columns_by_row_id, +}; use crate::index::{PreFilter, vector::VectorIndex}; use arrow::compute::concat_batches; use arrow_arith::numeric::sub; -use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array}; +use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array, cast::AsArray}; use arrow_schema::DataType; use async_trait::async_trait; use datafusion::error::{DataFusionError, Result as DataFusionResult}; @@ -911,7 +913,8 @@ impl IVFIndex { residual, scratch, )?; - Ok(batch) + // Emit covering columns with the per-partition result (batch path; no re-fetch). + Self::append_covering(batch, &part_entry.storage) } #[allow(clippy::too_many_arguments)] @@ -920,6 +923,7 @@ impl IVFIndex { use_residual_scratch: bool, prepared: PreparedPartitionSearch, heap: &mut BinaryHeap>, + covering_map: &mut HashMap, scratch: &mut QueryScratch, metrics: &dyn MetricsCollector, ) -> Result<()> { @@ -953,6 +957,8 @@ impl IVFIndex { let param = (&query).into(); let refine_factor = query.refine_factor.unwrap_or(1) as usize; let k = query.k * refine_factor; + // Accumulate this partition's contribution to the top-k heap first, so the heap + // membership consulted below already reflects it. part_entry.index.accumulate_topk_with_scratch( query.key, k, @@ -963,7 +969,70 @@ impl IVFIndex { residual, scratch, metrics, - ) + )?; + // Keep covering O(k): rather than buffering every probed partition's covering + // (O(nprobe * partition_size)) though only the k heap winners are emitted, extract + // only the covering rows for the current heap survivors -- deep-copied via + // `take_row` so this partition's storage can be released -- then prune the map to + // the heap's membership so rows evicted by this partition drop their covering. + // Every heap survivor's covering is present: a row is a survivor only if it was in + // the top-k when its own partition was processed (the heap never re-admits an + // evicted row), so it was extracted then -- or it is new from this partition. + // The lookup scans the partition's row ids against the (<= k)-sized needed set + // with an early exit, instead of building a partition-sized side map per probe. + if let Some(covering) = part_entry.storage.covering_batch()? { + let src_rowids = covering + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("covering batch missing row id".to_string()))? + .as_primitive::(); + let heap_ids: HashSet = heap.iter().map(|node| node.id).collect(); + let mut needed: HashSet = heap_ids + .iter() + .filter(|id| !covering_map.contains_key(id)) + .copied() + .collect(); + if !needed.is_empty() { + let mut positions: Vec = Vec::with_capacity(needed.len()); + let mut ids: Vec = Vec::with_capacity(needed.len()); + for (i, rid) in src_rowids.values().iter().enumerate() { + if needed.remove(rid) { + positions.push(i as u32); + ids.push(*rid); + if needed.is_empty() { + break; + } + } + } + if !positions.is_empty() { + // One take for this partition's whole contribution, not one per row: + // a per-row take ran an arrow take plus a `RecordBatch` schema + // validation to move a single cell, up to nprobe * k times per query. + let taken = arrow::compute::take_record_batch( + &covering, + &arrow_array::UInt32Array::from(positions), + )?; + // Each entry must OWN its row. A bare `taken.slice(offset, 1)` shares + // `taken`'s buffers, so one surviving winner would pin that whole + // partition's take batch -- with effective k = k * refine_factor, + // retained covering data becomes O(k^2) rows, not the O(k) this map + // exists to guarantee. `shrink_to_fit` deep-copies the single row, so + // both `taken` and the partition storage are released on drop. + for (offset, rid) in ids.into_iter().enumerate() { + covering_map.insert(rid, taken.slice(offset, 1).shrink_to_fit()?); + } + } + } + covering_map.retain(|id, _| heap_ids.contains(id)); + // Invariant: the map holds covering for exactly the heap's distinct row ids -- + // never more (that would break the O(k) bound) and never fewer (a survivor + // would be missing its covering at emit time). + debug_assert_eq!( + covering_map.len(), + heap_ids.len(), + "covering map must track exactly the heap's survivors (O(k))" + ); + } + Ok(()) } fn query_context_for_scratch<'a>( @@ -991,14 +1060,120 @@ impl IVFIndex { } } - fn global_heap_to_batch(heap: BinaryHeap>) -> Result { + fn global_heap_to_batch( + heap: BinaryHeap>, + // `row_id -> 1-row [_rowid, ]` for exactly the heap survivors, + // kept O(k) by `accumulate_prepared_partition_search`. + covering_map: &HashMap, + // `[_rowid, ]` schema for the index's covering columns, + // or None for an ordinary index. This is a stable per-index property, so + // it -- not `covering_map.is_empty()` -- decides whether to emit the wider + // covered schema. That keeps the emitted schema equal to the schema the + // exec declares even when zero partitions were searched (heap empty). + covering_schema: Option<&arrow_schema::Schema>, + ) -> Result { let (row_ids, dists): (Vec<_>, Vec<_>) = heap.into_iter().map(|r| (r.id, r.dist.0)).unzip(); + let dist_arr: ArrayRef = Arc::new(Float32Array::from(dists)); + let row_id_arr: ArrayRef = Arc::new(UInt64Array::from(row_ids)); + let Some(covering_schema) = covering_schema else { + // Ordinary index: `[_distance, _rowid]`. + return Ok(RecordBatch::try_new( + VECTOR_RESULT_SCHEMA.clone(), + vec![dist_arr, row_id_arr], + )?); + }; + // Covered index: emit `[_distance, _rowid, ]`. + let mut fields: Vec = VECTOR_RESULT_SCHEMA.fields().to_vec(); + let mut columns: Vec = vec![dist_arr, row_id_arr.clone()]; + if covering_map.is_empty() { + // No survivors (heap empty / zero partitions searched). Emit the covering + // columns as null arrays of the (zero) row count so the schema still matches + // the declared covered schema. + for field in covering_schema.fields() { + if field.name() == ROW_ID { + continue; + } + fields.push(field.clone()); + columns.push(arrow_array::new_null_array( + field.data_type(), + row_id_arr.len(), + )); + } + } else { + // Append the included columns for the final row ids, gathered from the O(k) + // side map of survivor covering rows captured in-flight during accumulate (no + // re-fetch from the base table). + let buf_schema = covering_map + .values() + .next() + .ok_or_else(|| { + Error::index( + "internal error: covering_map was empty despite the non-empty check \ + above" + .to_string(), + ) + })? + .schema(); + let combined = concat_batches(&buf_schema, covering_map.values())?; + let row_id_u64 = row_id_arr.as_primitive::(); + let included = gather_covering_columns_by_row_id(&combined, row_id_u64)?; + for (field, array) in included { + fields.push(field); + columns.push(array); + } + } Ok(RecordBatch::try_new( - VECTOR_RESULT_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(dists)), - Arc::new(UInt64Array::from(row_ids)), - ], + Arc::new(arrow_schema::Schema::new(fields)), + columns, + )?) + } + + /// Append the index's included/covering columns to a per-partition search + /// result (`[_distance, _rowid]`), gathered from the live partition storage. + /// No-op when the index has no covering columns. + /// Append the index's included/covering columns to a per-partition search + /// result (`[_distance, _rowid]`), gathered from the live partition storage. + /// No-op when the index has no covering columns. + /// The schema search results carry: `[_distance, _rowid]` plus any covering + /// ("included") columns the storage holds, in storage order. Used to declare + /// stream schemas that stay consistent with the (possibly widened) batches. + fn covered_result_schema( + covering_schema: Option<&arrow_schema::Schema>, + ) -> arrow_schema::SchemaRef { + let Some(covering_schema) = covering_schema else { + return VECTOR_RESULT_SCHEMA.clone(); + }; + let mut fields: Vec = VECTOR_RESULT_SCHEMA.fields().to_vec(); + for field in covering_schema.fields() { + if field.name() == ROW_ID { + continue; + } + fields.push(field.clone()); + } + Arc::new(arrow_schema::Schema::new(fields)) + } + + fn append_covering(batch: RecordBatch, storage: &Q::Storage) -> Result { + let Some(covering) = storage.covering_batch()? else { + return Ok(batch); + }; + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("search result missing row id".to_string()))? + .as_primitive::(); + let included = gather_covering_columns_by_row_id(&covering, row_ids)?; + if included.is_empty() { + return Ok(batch); + } + let mut fields: Vec = batch.schema().fields().to_vec(); + let mut columns: Vec = batch.columns().to_vec(); + for (field, array) in included { + fields.push(field); + columns.push(array); + } + Ok(RecordBatch::try_new( + Arc::new(arrow_schema::Schema::new(fields)), + columns, )?) } @@ -1366,6 +1541,24 @@ impl IVFIndex { self.storage.load_partition(partition_id, io_stats).await } + /// Names of this index's covering ("included") columns, in storage order, or + /// empty if the index has none. Used by remap/rebuild paths to re-project the + /// covering columns so they survive into the rewritten storage. + pub(crate) fn covering_column_names(&self) -> Result> { + Ok(self + .storage + .covering_schema()? + .map(|schema| { + schema + .fields() + .iter() + .filter(|f| f.name() != ROW_ID) + .map(|f| f.name().to_string()) + .collect() + }) + .unwrap_or_default()) + } + /// preprocess the query vector given the partition id. /// /// Internal API with no stability guarantees. @@ -1596,6 +1789,9 @@ impl VectorIndex for IVFInd scratch, ) })?; + // Emit any covering columns (parallel branch analog of the batch + // path's append_covering; no-op for ordinary indexes). + let batch = Self::append_covering(batch, &part_entry.storage)?; Result::Ok((batch, local_metrics)) }) .await?; @@ -1714,8 +1910,15 @@ impl VectorIndex for IVFInd let use_residual_scratch = self.use_residual_scratch; let search_metrics = metrics.clone(); let scratch_pool = self.scratch_pool.clone(); + // Stable per-index covering schema, so the emitted schema matches the + // covered schema even if zero partitions are searched (heap empty). + let covering_schema = self.storage.covering_schema()?; let batch = spawn_cpu(move || -> DataFusionResult { let mut heap = BinaryHeap::with_capacity(heap_capacity); + // O(k) side map of survivor covering rows (`row_id -> 1-row [_rowid, + // included...]`) captured in-flight; used to emit covering columns with the + // merged result (empty for ordinary indexes). + let mut covering_map: HashMap = HashMap::new(); scratch_pool.with_scratch(|scratch| -> DataFusionResult<()> { for prepared in prepared { Self::accumulate_prepared_partition_search( @@ -1723,6 +1926,7 @@ impl VectorIndex for IVFInd use_residual_scratch, prepared, &mut heap, + &mut covering_map, scratch, search_metrics.as_ref(), ) @@ -1730,12 +1934,16 @@ impl VectorIndex for IVFInd } Ok(()) })?; - Self::global_heap_to_batch(heap).map_err(DataFusionError::from) + Self::global_heap_to_batch(heap, &covering_map, covering_schema.as_deref()) + .map_err(DataFusionError::from) }) .await?; + // Schema may be wider than VECTOR_RESULT_SCHEMA when covering columns + // are emitted; take it from the produced batch so they stay consistent. + let result_schema = batch.schema(); return Ok(Box::pin(RecordBatchStreamAdapter::new( - VECTOR_RESULT_SCHEMA.clone(), + result_schema, stream::once(async move { Ok(batch) }), ))); } @@ -1922,8 +2130,12 @@ impl VectorIndex for IVFInd } }); + // The per-partition batches are widened with the storage's covering columns + // (`append_covering`), so declare the matching schema -- the global-heap + // branch above already emits its batch's own (covered) schema. + let result_schema = Self::covered_result_schema(self.storage.covering_schema()?.as_deref()); Ok(Box::pin(RecordBatchStreamAdapter::new( - VECTOR_RESULT_SCHEMA.clone(), + result_schema, ReceiverStream::new(batch_rx), ))) } @@ -2120,12 +2332,15 @@ mod tests { use lance_index::vector::storage::VectorStore; use lance_index::vector::v3::subindex::IvfSubIndex; - use crate::dataset::{InsertBuilder, UpdateBuilder, WriteMode, WriteParams}; + use crate::dataset::{ + InsertBuilder, NewColumnTransform, UpdateBuilder, WriteMode, WriteParams, + }; use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; use crate::index::vector::ivf::v2::{ IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, PartitionEntry, }; + use crate::index::vector::utils::gather_covering_columns_by_row_id; use crate::utils::test::copy_test_data_to_tmp; use crate::{ Dataset, @@ -2945,6 +3160,1106 @@ mod tests { } } + /// End-to-end test for index-included ("covering") columns : an + /// `covering_columns` request on the build params must survive the full + /// build -> shuffle -> persist -> reopen path and land in the IVF_PQ + /// partition storage, so covered queries can avoid a take. + #[tokio::test] + async fn test_ivf_pq_covering_columns_roundtrip() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + // `generate_test_dataset` produces an `id` (UInt64) column besides `vector`. + let (mut dataset, _) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "included column 'id' should be co-located in IVF_PQ partition storage" + ); + + // The covering columns are also declared generically on IndexMetadata + // (by field id), so the read path can discover them for any index type. + let id_field_id = dataset.schema().field("id").unwrap().id; + let indices = dataset.load_indices().await.unwrap(); + let idx = indices + .iter() + .find(|i| i.name == INDEX_NAME) + .expect("index should exist"); + assert_eq!( + idx.covering_fields, + vec![id_field_id], + "IndexMetadata.covering_fields should record the covered column's field id" + ); + + // The REAL producer must raise the reader+writer fence, not just the + // doctored-metadata commit that upstream's + // `test_covering_commit_fences_the_table_with_a_feature_flag` exercises: a + // pre-covering build that opened this dataset would select the index by + // membership of `fields` and answer a query on the carried column with + // an index keyed on another one. + use lance_table::feature_flags::FLAG_COVERED_INDEX_METADATA; + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_COVERED_INDEX_METADATA, + 0, + "a covered create must raise the reader fence" + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_COVERED_INDEX_METADATA, + 0, + "a covered create must raise the writer fence" + ); + } + + /// A multivector column stores one entry per sub-vector (all sharing the source + /// row id). The covering build must preserve every sub-vector -- the stale-row + /// dedup must not collapse legitimate repeated row ids, or recall silently drops. + #[tokio::test] + async fn test_ivf_pq_covered_multivector_preserves_all_subvectors() { + const INDEX_NAME: &str = "vector_idx"; + const SUBVECS_PER_ROW: usize = 3; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, _) = + generate_multivec_test_dataset::(test_uri, 0.0..1.0).await; + + // Multivector requires cosine. Cover the `id` column. + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::Cosine, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let mut total_stored = 0usize; + let mut has_id = false; + for p in 0..ctx.num_partitions() { + let storage = ctx.ivf().load_partition_storage(p, None).await.unwrap(); + total_stored += storage.batch().num_rows(); + has_id |= storage.batch().column_by_name("id").is_some(); + } + assert!( + has_id, + "multivector covered storage should carry covering column 'id'" + ); + assert_eq!( + total_stored, + NUM_ROWS * SUBVECS_PER_ROW, + "covering multivector build must keep every sub-vector, not collapse per row id" + ); + } + + /// Read-side payoff for multivector: a covered projection is carried through + /// `MultivectorScoringExec` (which re-groups sub-vectors back to rows), so no + /// `TakeExec` against the base table is needed. + #[tokio::test] + async fn test_ivf_pq_covered_multivector_projection_skips_take() { + use arrow_array::types::UInt64Type; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let (mut dataset, vectors) = + generate_multivec_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::Cosine, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let query = vectors.value(0); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.minimum_nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered multivector projection ['id'] should skip TakeExec; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted through multivector scoring") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(batch.num_rows() > 0, "query should return rows"); + // Single-fragment, step-id dataset => id == row offset == _rowid, so a correctly + // carried covered value equals the row id for every returned row (payload + // re-attached to the right row through scoring, not misaligned or stale). + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + row_ids.value(i), + "covered 'id' must stay row-aligned through multivector scoring" + ); + } + } + + /// A dotted/nested include path cannot be covered (the covering gather projects only + /// top-level names); reject it at build time. + #[tokio::test] + async fn test_ivf_pq_rejects_dotted_include_column() { + let test_dir = TempStrDir::default(); + let (mut dataset, _) = + generate_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id.sub".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("a dotted include column must be rejected"); + assert!( + err.to_string().contains("nested/dotted"), + "expected a nested/dotted rejection, got: {err}" + ); + } + + /// Covering names that collide with index storage (the indexed vector column itself, + /// a reserved storage column, or a duplicate) must be rejected at build time. + #[tokio::test] + async fn test_ivf_rejects_reserved_and_duplicate_covering_columns() { + let test_dir = TempStrDir::default(); + let (mut dataset, _) = + generate_test_dataset::(test_dir.as_str(), 0.0..1.0).await; + let build = |cols: Vec| { + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(cols); + params + }; + + // The indexed vector column itself. + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec!["vector".to_string()]), + true, + ) + .await + .expect_err("covering the indexed vector column must be rejected"); + assert!( + err.to_string().contains("indexed vector column itself"), + "got: {err}" + ); + + // Reserved storage/transform column names. The reserved check runs before the + // schema-existence check, so these are rejected even though the dataset lacks them. + // The RaBitQ code/factor columns and the IVF partition transform's `__centroid_dist` + // are internal to the build pipeline: covering one would advertise it in + // `covering_fields` while the storage's `covering_field_indices` drops it, so a + // covered query would declare a column storage never emits. + for internal in [ + "__pq_code", + "__sq_code", + "__ivf_part_id", + "__centroid_dist", + "__ex_codes", + "__blocked_ex_codes", + "__add_factors", + "__scale_factors", + "__error_factors", + "__add_factors_ex", + "__scale_factors_ex", + ] { + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec![internal.to_string()]), + true, + ) + .await + .expect_err("covering a reserved storage/transform name must be rejected"); + assert!( + err.to_string().contains("reserved index storage"), + "internal name '{internal}' must be reserved, got: {err}" + ); + } + + // Duplicate covering columns. + let err = dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &build(vec!["id".to_string(), "id".to_string()]), + true, + ) + .await + .expect_err("duplicate covering columns must be rejected"); + assert!(err.to_string().contains("duplicate"), "got: {err}"); + } + + /// A blob column stores out-of-line descriptors, not inline data, so it cannot be + /// covered by a vector index; reject it at build time. + #[tokio::test] + async fn test_ivf_flat_rejects_blob_include_column() { + use arrow_array::{ + Float32Array, Int32Array, LargeBinaryArray, RecordBatch, RecordBatchIterator, + }; + use lance_arrow::BLOB_META_KEY; + use lance_file::version::LanceFileVersion; + use std::collections::HashMap; + + let dim = 4i32; + let n = 32usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( + BLOB_META_KEY.to_string(), + "true".to_string(), + )])), + ])); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n as i32 * dim).map(|v| v as f32).collect::>()), + dim, + ) + .unwrap(); + let blobs: Vec> = (0..n).map(|_| Some(b"x".as_slice())).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new(vectors), + Arc::new(LargeBinaryArray::from(blobs)), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://c1-blob-include", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }), + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["blobs".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("a blob include column must be rejected"); + assert!( + err.to_string().contains("blob column"), + "expected a blob rejection, got: {err}" + ); + } + + /// Read-side payoff : a vector query projecting only a covered column + /// is satisfied from the index — no `TakeExec` against the base table. + #[tokio::test] + async fn test_ivf_pq_covered_projection_skips_take() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + + // Take elision is worthless if the covered search returns the wrong + // neighbors, so also gate on recall: compare the covered result's row ids + // against brute-force ground truth (use_index(false)). All 4 partitions + // are probed, so 0.5 sits far below IVF_PQ's real recall on this data. + let returned: HashSet = batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered IVF_PQ recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// The most likely real query shape: a projection mixing a covered column + /// with an uncovered one. `filtered_read.rs` re-subtracts the projection + /// against the covered stream's schema, so `id` must not be re-fetched + /// while `extra` -- a column the index does not carry -- still goes + /// through a `TakeExec`. Every other covered test projects only-covered + /// or only-uncovered columns and would not catch a regression here. + #[tokio::test] + async fn test_ivf_pq_covered_partial_projection() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + // 16 sub-vectors, not 4: at DIM=32 that is 2 dimensions per codebook instead of + // 8. This fixture is uniform-random, so the true top-10 is full of near-ties, and + // at 4 sub-vectors the PQ error is large enough that NEON-vs-AVX rounding reorders + // them -- the recall gate below read 0.4 on aarch64 while passing on x86. The gate + // is a sanity check on take elision, not a measurement of quantizer accuracy. + let mut params = VectorIndexParams::ivf_pq(4, 8, 16, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A column the index does not cover; the take path must still fetch it. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![( + "extra".to_string(), + "CAST(id AS BIGINT) + 1000".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "extra"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "uncovered column 'extra' must still require a TakeExec; plan was:\n{plan}" + ); + assert!( + plan.contains("projection=[extra]"), + "the take must fetch only the uncovered column, not re-fetch the \ + already-covered 'id'; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + let ids = batch["id"].as_primitive::(); + let extras = batch["extra"].as_primitive::(); + for (id, extra) in ids.values().iter().zip(extras.values().iter()) { + assert_eq!( + *extra, + *id as i64 + 1000, + "uncovered 'extra' must match the base table, not a stale/misaligned value" + ); + } + + // Take elision only matters if the covered result is correct. + let returned: HashSet = batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered IVF_PQ recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// Two segments of one logical index disagreeing on `covering_fields` is + /// exactly the corruption the read path cannot tolerate: `knn.rs` and + /// `scanner.rs` both derive the exec's declared output schema from a + /// single segment and assume every sibling of the same logical index + /// agrees. Reproduce the failure with public APIs alone: a covered + /// index, then a fresh batch of fragments indexed *without* covering (a + /// caller who forgot `covering_columns`, or a differently configured + /// distributed build), committed alongside the still-covered original + /// segment. Both segments are keyed on the same field with `keyed == 1`, + /// so every other commit-time check passes; only the `covering_fields` + /// agreement check added here catches it. + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_covering_fields_disagreement() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, _vectors) = + generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut covered_params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + covered_params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &covered_params, + true, + ) + .await + .unwrap(); + + let original_fragment_ids: HashSet = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect(); + // PQ training on the new segment needs >= 256 rows (2^8 codes). + append_dataset::(&mut dataset, 300, 0.0..1.0).await; + let new_fragment_ids: Vec = dataset + .get_fragments() + .into_iter() + .map(|f| f.id() as u32) + .filter(|id| !original_fragment_ids.contains(id)) + .collect(); + assert!(!new_fragment_ids.is_empty(), "append should add fragments"); + + // Build a segment for just the new fragments, without covering. + let uncovered_params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + let uncovered_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &uncovered_params) + .name(INDEX_NAME.to_string()) + .fragments(new_fragment_ids) + .replace(true) + .execute_uncommitted() + .await + .unwrap(); + assert!( + uncovered_segment.covering_fields.is_empty(), + "the new segment must genuinely be uncovered for this repro" + ); + + // Committing it alongside the still-covered original segment must be + // rejected, not silently accepted into one logical index that + // disagrees with itself on `covering_fields`. + let err = dataset + .commit_existing_index_segments(INDEX_NAME, "vector", vec![uncovered_segment]) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("a logical index cannot mix declarations"), + "unexpected error: {err}" + ); + } + + /// Same payoff but forcing the BATCH/late-search path: `k` larger than one + /// partition makes the initial `minimum_nprobes` sweep under-fill, so + /// `late_search` (the run_prepared batch path) fires. Covering must hold there too. + #[tokio::test] + async fn test_ivf_pq_covered_projection_batch_path() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + // k >> one partition's rows (NUM_ROWS=512 / 4 parts) forces late expansion. + scan.nearest("vector", q, 500).unwrap(); + scan.minimum_nprobes(1); + scan.maximum_nprobes(4); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec (batch path); plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + assert!(batch.num_rows() > 100, "late search should have expanded"); + } + + /// A covered query combined with a selective prefilter (scalar index + + /// `prefilter`) must still skip the TakeExec and return correct, filtered + /// rows. Exercises the covering read path under a scalar-index prefilter. + #[tokio::test] + async fn test_ivf_pq_covered_with_scalar_prefilter() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A scalar index on `id` turns the prefilter into a bounded AllowList. + let scalar_params = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_btree".to_string()), + &scalar_params, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.filter("id < 5").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec even with a prefilter; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + // Every id in [0, 5) is present in the dataset, so the prefilter must + // return exactly 5 rows -- not zero, which would make the loop below + // vacuously true regardless of whether filtering actually worked. + assert_eq!(batch.num_rows(), 5); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + for v in ids.values() { + assert!( + *v < 5, + "all returned rows must satisfy the prefilter id < 5" + ); + } + } + + /// A covered index must keep its covering columns through optimize/merge. + /// After appending rows and merging them into the index, the merged + /// auxiliary storage must still carry the included column, and a covered + /// projection must still skip the take. The covering columns are threaded + /// through the incremental optimize pipeline -- the unindexed-fragment + /// shuffle, the partition-split reshuffle, and the partition-join + /// reassignment -- as passenger columns (re-gathered by row id). + #[tokio::test] + async fn test_ivf_pq_covered_survives_optimize() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Append rows (unindexed fragments) and merge them into the index. + append_dataset::(&mut dataset, NUM_ROWS, 0.0..1.0).await; + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + // The merged partition storage must still carry the covering column. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "merged storage should still carry covering column 'id'" + ); + + // And a covered projection is still answered from the index (no take). + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after optimize; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + } + + /// Covering must survive a *retrain* optimize. Retrain rebuilds the storage from + /// scratch, so the covering column must be re-materialized -- otherwise the fresh + /// files omit it while the committed metadata still advertises `covering_fields`, + /// and a covered projection expects a column the storage cannot emit. + #[tokio::test] + async fn test_ivf_pq_covered_survives_retrain() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + append_dataset::(&mut dataset, NUM_ROWS, 0.0..1.0).await; + dataset + .optimize_indices(&OptimizeOptions::retrain()) + .await + .unwrap(); + + // The retrained storage must re-materialize the covering column. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "retrained storage should re-materialize covering column 'id'" + ); + + // And a covered projection is still answered from the index (no take). + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after retrain; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + } + + /// Index storage that physically carries covering columns while the manifest + /// declares none (`covering_fields` empty) must still be queryable: the extra + /// columns are dropped from search batches (with a warning) instead of emitting + /// batches wider than the plan's declared `[_distance, _rowid]` schema. This is + /// the legacy-tolerance contract: a stable-format index file with an unexpected + /// extra column used to be projected down with a warning, never a query failure. + #[tokio::test] + async fn test_undeclared_storage_covering_is_dropped_not_fatal() { + use crate::dataset::WriteDestination; + use crate::dataset::transaction::Operation; + + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + // Well-separated clusters, not uniform-random vectors: a query near one + // cluster must land far outside the early-pruning heuristic's threshold + // (`early_pruning`/`adjust_probes` in `io/exec/knn.rs`) for every other + // cluster's centroid, so `late_search` gets a real range of partitions + // left to search instead of the heuristic folding all of them into the + // early, synchronous phase. Uniform-random vectors near the dataset's + // overall centroid are roughly equidistant from every partition centroid, + // so every partition is searched up front and `late_search` never runs -- + // silently defeating the "hard case" below (a defect recorded once + // already on this project: a covered-ANN test with too few effectively- + // reachable partitions passes without ever exercising the path it claims + // to cover). + let mut rng = StdRng::seed_from_u64(7); + let mut ids = Vec::with_capacity(TOTAL); + let mut values = Vec::with_capacity(TOTAL * DIMS); + for cluster in 0..NUM_CLUSTERS { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + ids.push((cluster * ROWS_PER_CLUSTER + row) as i32); + for dim in 0..DIMS { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + (rng.random::() - 0.5) * 0.02); + } + } + } + let ids_arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(ids)); + let vectors: ArrayRef = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIMS as i32) + .unwrap(), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids_arr, vectors]).unwrap(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Re-commit the index metadata with the covering declaration cleared while the + // storage keeps the payload column -- the state a pre-covering writer (which + // drops the unknown `covering_fields` proto field) or a legacy index file with + // a stray extra column would produce. + let mut cleared = dataset.load_indices_by_name("vector_idx").await.unwrap()[0].clone(); + assert!(!cleared.covering_fields.is_empty()); + cleared.covering_fields = Vec::new(); + let read_version = dataset.manifest.version; + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::CreateIndex { + new_indices: vec![cleared], + removed_indices: vec![], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // Query the LAST cluster's center: its rows' ids (`3*ROWS_PER_CLUSTER..`) + // never satisfy the `id < 3` filter used below, so the "hard case" prefilter + // match (ids 0/1/2, in the first cluster) is guaranteed to be entirely + // unfound by the early search of this, the nearest, cluster/partition -- + // exercising the not-found shortcut instead of incidentally short-circuiting + // because the early search already found every matching row. + let mut q_values = vec![0.0f32; DIMS]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let q = Float32Array::from(q_values); + + // The index is no longer covering, so the query takes `id` from the base + // table; the storage's undeclared payload column must be silently dropped + // from the search batches, not fail the query. + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + assert!(batch.column_by_name("id").is_some()); + + // The hard case: a bounded selective prefilter mixes the plan's 2-column + // not-found shortcut batch with search batches -- if the search batches still + // carried the undeclared payload column, the widths would disagree where the + // stream is concatenated and the query would fail. + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.filter("id < 3").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + 3, + "all prefilter-matched rows must come back" + ); + } + + /// A covered index that searches zero partitions (empty heap) + /// must still emit the covered schema `[_distance, _rowid, ]`, not + /// the bare `[_distance, _rowid]` -- otherwise the produced batch mismatches + /// the wider schema the exec declares from `covering_fields`. The decision is + /// driven by the stable per-index covering schema, not the gathered `covering` batch. + #[test] + fn test_global_heap_to_batch_covered_empty_emits_covered_schema() { + use arrow_schema::{DataType, Field, Schema}; + let covering = Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new("id", DataType::UInt64, true), + ]); + // Empty heap + empty covered_buf, but the index IS covered. + let batch = IvfPq::global_heap_to_batch( + std::collections::BinaryHeap::new(), + &HashMap::new(), + Some(&covering), + ) + .unwrap(); + assert_eq!(batch.num_rows(), 0); + assert_eq!( + batch.num_columns(), + 3, + "covered empty result must keep the covering column" + ); + assert!(batch.column_by_name("id").is_some()); + + // Ordinary (non-covered) index: bare `[_distance, _rowid]`. + let plain = + IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), &HashMap::new(), None) + .unwrap(); + assert_eq!(plain.num_columns(), 2); + } + + /// A result row id that is absent from the covering buffer is an invariant + /// break (every heap winner comes from a searched partition whose covering + /// batch was captured). It must surface as an error -- never as row 0's + /// covering values silently attached to an unrelated row. + #[test] + fn test_gather_covering_by_rowid_errors_on_missing_rowid() { + use arrow_schema::{DataType, Field, Schema}; + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new("id", DataType::UInt64, true), + ])); + let source = RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from(vec![1, 2, 3])), + Arc::new(UInt64Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + + // Present row ids gather fine. + let ok = + gather_covering_columns_by_row_id(&source, &UInt64Array::from(vec![3, 1])).unwrap(); + assert_eq!(ok.len(), 1); + let (_, values) = &ok[0]; + let values = values.as_primitive::(); + assert_eq!(values.values(), &[30, 10]); + + // Row id 99 is absent from the covering source: must be an error. + let err = gather_covering_columns_by_row_id(&source, &UInt64Array::from(vec![2, 99])) + .expect_err("missing row id must error, not return unrelated values"); + assert!(err.to_string().contains("99"), "got: {err}"); + } + + /// A covered index must keep its covering columns through compaction, which + /// rewrites fragments and REMAPS the index's row ids. Delete rows to force a + /// compaction+remap; the remapped storage must still carry the covering + /// column and covered projections must still skip the take. + #[tokio::test] + async fn test_ivf_pq_covered_survives_compaction() { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows so compaction rewrites fragments and remaps the index. + dataset.delete("id < 100").await.unwrap(); + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + + // The remapped partition storage must still carry the covering column. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + assert!( + storage.batch().column_by_name("id").is_some(), + "remapped storage should keep covering column 'id'" + ); + + // Covered projection still skips the take, and deleted rows are gone. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after compaction; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + // Plenty of non-deleted rows remain for k=10 to fill; guard against an + // empty result, which would make the loop below vacuously true. + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered projection should return id") + .as_primitive::(); + for v in ids.values() { + assert!(*v >= 100, "deleted rows (id < 100) must not be returned"); + } + } + + /// Fix 2: under query parallelism > 1 the parallel search branch (which uses + /// `search_in_partition`, not the heap merge) must also emit covering + /// columns. On a multi-core runner this exercises the parallel path; if the + /// session resolves parallelism to 1 it still passes via the sequential path. + #[tokio::test] + async fn test_ivf_pq_covered_projection_parallel() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.minimum_nprobes(4); // search several partitions... + scan.query_parallelism(4); // ...in parallel, forcing the parallel branch + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + } + + /// Fix 3: a covered index with unindexed fragments (rows appended but not + /// optimized) must not panic on a non-fast-search query. The flat search + /// over the unindexed rows now projects the covering columns so the union + /// with the index result succeeds. + #[tokio::test] + async fn test_ivf_pq_covered_with_unindexed_fragments() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Append rows WITHOUT optimizing -> unindexed fragments; a (non-fast) + // query then goes through the index + flat-search combine path. + append_dataset::(&mut dataset, 100, 0.0..1.0).await; + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert!(batch.column_by_name("id").is_some()); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + } + async fn shrink_smallest_partition( dataset: &mut Dataset, index_name: &str, diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index 012d8d3a703..b105cdfbefd 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -23,6 +23,52 @@ use tokio::sync::Mutex; use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder, row_offsets_to_row_addresses}; use crate::{Error, Result}; +use lance_core::ROW_ID; + +/// Row-align the non-`_rowid` columns of `source` (a `[_rowid, ]` batch) +/// to `target_rowids`, gathered by row id. Every target row id must be present in +/// `source`; a miss returns an error, because silently gathering an unrelated +/// row's values would corrupt the covered result. Returns the taken columns paired +/// with their fields, in `source` column order, with `_rowid` excluded. +/// +/// Shared by the covered-search read path (regrouping heap winners into +/// `[_distance, _rowid, ]`) and the reassignment build path +/// (re-attaching covering columns to rows regrouped by the partition join). +pub(super) fn gather_covering_columns_by_row_id( + source: &RecordBatch, + target_rowids: &arrow_array::UInt64Array, +) -> Result> { + let src_rowids = source + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("covering source missing row id".to_string()))? + .as_primitive::(); + let mut pos: std::collections::HashMap = + std::collections::HashMap::with_capacity(src_rowids.len()); + for (i, rid) in src_rowids.values().iter().enumerate() { + pos.insert(*rid, i as u32); + } + let take_idx: arrow_array::UInt32Array = target_rowids + .values() + .iter() + .map(|rid| { + pos.get(rid).copied().ok_or_else(|| { + Error::internal(format!( + "row id {rid} missing from covering source during row-id gather" + )) + }) + }) + .collect::>>()? + .into(); + let mut out = Vec::with_capacity(source.num_columns().saturating_sub(1)); + for (i, field) in source.schema().fields().iter().enumerate() { + if field.name() == ROW_ID { + continue; + } + let taken = arrow::compute::take(source.column(i), &take_idx, None)?; + out.push((field.clone(), taken)); + } + Ok(out) +} /// Helper function to extract a column from a RecordBatch, supporting nested field paths. /// diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 7a334198610..2bdc556abd9 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -62,7 +62,7 @@ use roaring::RoaringBitmap; use tokio::sync::Notify; use uuid::Uuid; -use crate::dataset::Dataset; +use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder}; use crate::index::DatasetIndexInternalExt; use crate::index::prefilter::{DatasetPreFilter, FilterLoader}; use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; @@ -1399,6 +1399,19 @@ pub struct ANNIvfSubIndexExec { /// Optional external row-address allow/block mask, combined with the /// prefilter using logical AND. external_mask: Option>, + /// Output schema: `[_distance, _rowid]` plus any included/covering columns + /// the index emits. + output_schema: SchemaRef, + + /// Whether this plan emits covering columns, resolved from the index + /// declaration at construction. Kept as its own field rather than re-derived + /// from the output schema's width at execution: width-based inference is + /// exact only while covering is the sole thing that widens the schema, and + /// the first unrelated extra column would silently arm per-query emitted-row + /// tracking and the covered recovery take on every plain query (or mask + /// covering if a column were dropped) with results staying correct -- no test + /// would catch it. + has_covered: bool, /// Datafusion Plan Properties properties: Arc, @@ -1420,8 +1433,46 @@ impl ANNIvfSubIndexExec { PART_ID_COLUMN ))); } + // Declare any included/covering columns recorded in the index manifest + // (`IndexMetadata.covering_fields`, by field id); the per-partition search + // emits them so a covered projection lets TakeExec skip the base-table + // fetch. Empty for ordinary indexes. Resolve id -> name -> arrow field so + // the declared schema matches the columns stored at build time exactly. + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + let included_ids = indices + .first() + .map(|idx| idx.covering_fields.clone()) + .unwrap_or_default(); + if !included_ids.is_empty() { + let lance_schema = dataset.schema(); + for id in &included_ids { + // The manifest is authoritative: a declared covering field that cannot + // be resolved means the index metadata and schema disagree. Silently + // dropping it would leave the declared output schema missing a column a + // covered projection expects -- TakeExec would then wrongly elide the + // base-table fetch for a column that is never emitted. Fail loudly + // instead (a covered column cannot be dropped while its index exists). + let lance_field = lance_schema + .fields + .iter() + .find(|field| field.id == *id) + .ok_or_else(|| { + Error::index(format!( + "ANNIvfSubIndexExec: index declares covering field id {id}, which is \ + not present as a top-level field in the current dataset schema; \ + index metadata and schema are inconsistent" + )) + })?; + // Convert just this field. `ArrowSchema::from(&lance_schema)` would give the + // same result -- it is exactly this conversion mapped over every field -- but + // it walks the whole table (including nested structs) on every covered plan. + fields.push(Arc::new(arrow_schema::Field::from(lance_field))); + } + } + let has_covered = !included_ids.is_empty(); + let output_schema = Arc::new(arrow_schema::Schema::new(fields)); let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(KNN_INDEX_SCHEMA.clone()), + EquivalenceProperties::new(output_schema.clone()), Partitioning::RoundRobinBatch(1), EmissionType::Final, Boundedness::Bounded, @@ -1434,6 +1485,8 @@ impl ANNIvfSubIndexExec { prefilter_source, overlay_block: None, external_mask: None, + output_schema, + has_covered, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1513,6 +1566,16 @@ struct ANNIvfEarlySearchResults { deltas_remaining: AtomicUsize, all_deltas_done: Notify, took_no_rows_shortcut: AtomicBool, + /// One entry per covered delta whose late search reached the no-rows shortcut: + /// that delta's per-segment ownership mask (`None` when the delta owned the + /// whole global mask). The exec-level covered recovery emits ONLY rows these + /// scopes select, so covered result sets stay identical to non-covered ones in + /// both directions: a delta that never armed (it skipped its late search, e.g. + /// probing fewer partitions than `minimum_nprobes`) would not have emitted its + /// not-found rows on the non-covered path either, so recovery must not emit + /// them for it -- and an empty list means no delta armed, so recovery emits + /// nothing at all. + covered_recovery_scopes: Mutex>>>, } impl ANNIvfEarlySearchResults { @@ -1524,6 +1587,7 @@ impl ANNIvfEarlySearchResults { deltas_remaining: AtomicUsize::new(deltas_remaining), all_deltas_done: Notify::new(), took_no_rows_shortcut: AtomicBool::new(false), + covered_recovery_scopes: Mutex::new(Vec::new()), } } @@ -1782,6 +1846,7 @@ impl ANNIvfSubIndexExec { state: Arc, target_partitions: usize, seg_mask: Option>, + has_covered: bool, ) -> impl Stream> { let stream = futures::stream::once(async move { let max_nprobes = query @@ -1826,6 +1891,29 @@ impl ANNIvfSubIndexExec { // In this case there are fewer than k results matching the prefilter so // just return the prefilter ids and don't bother searching any further + // The shortcut batch emits only [_distance, _rowid]; a covering + // index's declared schema carries more columns, and their values + // cannot be fabricated here. Emit nothing instead: the + // end-of-stream covered recovery (`covered_not_found_batch`) + // returns every unemitted prefilter row with INFINITY distance + // and its covering columns via a take bounded by the prefilter + // size -- mirroring the non-covered shortcut's semantics without + // searching (a prefilter row with no index entry, e.g. a null + // vector, could otherwise never be found and would drive the + // search through every partition). + if has_covered { + // Arm recovery with THIS delta's ownership scope: recovery may + // emit only rows the arming deltas own, or it diverges from the + // non-covered result set whenever a sibling delta skipped its + // late search without arming (see `covered_recovery_scopes`). + state + .covered_recovery_scopes + .lock() + .unwrap() + .push(seg_mask.clone()); + return futures::stream::empty().boxed(); + } + // This next if check should be true, because we wouldn't get max_results otherwise if let Some(iter_addrs) = prefilter_mask.iter_addrs() { // Emit the prefilter rows that the partition search did not reach. @@ -2029,6 +2117,93 @@ impl ANNIvfSubIndexExec { .buffered(query_parallelism) .boxed() } + + /// Recover rows that match a bounded, selective prefilter but have no index entry + /// (their vector is null or an empty multivector), which the covered search never + /// returns. The non-covered path recovers them via its not-found shortcut with + /// `INFINITY` distance; a covered index must do the same but also supply the covering + /// columns, which it fetches from the base table for these few missing rows. Only + /// rows selected by an arming delta's ownership scope in `scopes` are recovered + /// (`None` = that delta owned the whole global mask): a delta that skipped its late + /// search without arming never emits its not-found rows on the non-covered path, so + /// recovering them here would make the covered result set diverge. Returns `None` + /// when there is nothing to recover (non-covered, unbounded/non-selective prefilter, + /// or every prefilter row was already emitted or unowned). + async fn covered_not_found_batch( + has_covered: bool, + prefilter: Arc, + dataset: Arc, + output_schema: SchemaRef, + k: usize, + emitted: Arc>, + scopes: Vec>>, + ) -> DataFusionResult> { + if !has_covered { + return Ok(None); + } + prefilter + .wait_for_ready() + .await + .map_err(|e| DataFusionError::Execution(format!("prefilter not ready: {e}")))?; + let mask = prefilter.mask(); + // Only a bounded, selective prefilter (at most k allowed rows) guarantees every + // matching row belongs in the result; an unbounded prefilter has no not-found set. + let Some(max_len) = mask.max_len() else { + return Ok(None); + }; + if max_len as usize > k { + return Ok(None); + } + let Some(addrs) = mask.iter_addrs() else { + return Ok(None); + }; + let not_found: Vec = { + let emitted = emitted.lock().unwrap(); + addrs + .map(u64::from) + .filter(|addr| !emitted.contains(*addr)) + .filter(|addr| { + // Mirror the non-covered shortcut's per-segment ownership rule: + // each armed delta may recover only the rows its segment owns. + scopes + .iter() + .any(|scope| scope.as_ref().is_none_or(|m| m.selected(*addr))) + }) + .collect() + }; + if not_found.is_empty() { + return Ok(None); + } + + // Fetch the covering columns (the fields after `[_distance, _rowid]`) from the + // base table for the missing rows; the take is bounded by `not_found.len()`. + let covered_names: Vec = output_schema + .fields() + .iter() + .skip(2) + .map(|field| field.name().clone()) + .collect(); + let projection = ProjectionRequest::from_columns(covered_names, dataset.schema()); + // `not_found` holds row ids (the `_rowid` space): stable logical ids on a + // stable-row-id dataset, physical addresses otherwise. `try_new_from_ids` + // resolves them through the row-id index when one exists and falls back to + // identity (id == address) when it does not, so it is correct in both cases -- + // unlike `try_new_from_addresses`, which would misread a stable id as + // `frag = id >> 32, offset = id` and take the wrong physical row (or error). + let covered = TakeBuilder::try_new_from_ids(dataset, not_found.clone(), projection)? + .execute() + .await?; + + let n = not_found.len(); + let mut columns: Vec = vec![ + Arc::new(Float32Array::from_value(f32::INFINITY, n)), + Arc::new(UInt64Array::from(not_found)), + ]; + columns.extend(covered.columns().iter().cloned()); + let batch = RecordBatch::try_new(output_schema, columns) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + Ok(Some(batch)) + } } impl ExecutionPlan for ANNIvfSubIndexExec { @@ -2037,7 +2212,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { } fn schema(&self) -> arrow_schema::SchemaRef { - KNN_INDEX_SCHEMA.clone() + self.output_schema.clone() } fn children(&self) -> Vec<&Arc> { @@ -2084,6 +2259,8 @@ impl ExecutionPlan for ANNIvfSubIndexExec { prefilter_source, overlay_block: self.overlay_block.clone(), external_mask: self.external_mask.clone(), + output_schema: self.output_schema.clone(), + has_covered: self.has_covered, properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -2102,6 +2279,10 @@ impl ExecutionPlan for ANNIvfSubIndexExec { ) -> DataFusionResult { let input_stream = self.input.execute(partition, context.clone())?; let schema = self.schema(); + // When the index emits covering columns, the late-search no-rows shortcut + // (which emits only [_distance, _rowid]) is disabled in favour of the + // end-of-stream covered recovery. Resolved at construction; see the field. + let has_covered = self.has_covered; let target_partitions = context.session_config().target_partitions(); let query = self.query.clone(); let ds = self.dataset.clone(); @@ -2202,6 +2383,22 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let state = Arc::new(ANNIvfEarlySearchResults::new(indices.len(), query.k)); + // Covered null-vector recovery: track which row ids the search emits, then after + // it finishes emit any prefilter-matched rows that had no index entry, with their + // covering columns taken from the base table (see `covered_not_found_batch`). + // Tracking is skipped when recovery provably cannot emit anything: without a + // prefilter source there is no allow-list, so `covered_not_found_batch` always + // returns `None` -- the per-batch lock and inserts would be pure overhead on + // the covered hot path. + let track_emitted = has_covered && !matches!(self.prefilter_source, PreFilterSource::None); + let emitted_row_ids = Arc::new(std::sync::Mutex::new(roaring::RoaringTreemap::new())); + let emitted_for_inspect = emitted_row_ids.clone(); + let recon_prefilter = pre_filter.clone(); + let recon_state = state.clone(); + let recon_ds = ds.clone(); + let recon_schema = schema.clone(); + let recon_k = query.k; + Ok(Box::pin(RecordBatchStreamAdapter::new( schema, per_index_stream @@ -2277,6 +2474,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state, target_partitions, seg_mask, + has_covered, ); DataFusionResult::Ok(early_search.chain(late_search).boxed()) } @@ -2285,6 +2483,42 @@ impl ExecutionPlan for ANNIvfSubIndexExec { // Each delta stream is split into an early and late search. The late search // will not start until the early search is complete across all deltas. .try_flatten_unordered(None) + .inspect_ok(move |batch| { + // Record emitted row ids so the reconciliation below can find the + // prefilter rows the search never returned (no index entry). + if track_emitted && let Some(col) = batch.column_by_name(ROW_ID) { + emitted_for_inspect + .lock() + .unwrap() + .extend(col.as_primitive::().values().iter().copied()); + } + }) + .chain( + stream::once(async move { + // Only recover when a covered late search actually took the + // no-rows shortcut; otherwise the non-covered path would not + // have emitted these rows either, and covered must match. + // Each armed delta contributes its ownership scope, and the + // recovery emits only rows those scopes select. + let scopes = std::mem::take( + &mut *recon_state.covered_recovery_scopes.lock().unwrap(), + ); + if scopes.is_empty() { + return Ok(None); + } + Self::covered_not_found_batch( + has_covered, + recon_prefilter, + recon_ds, + recon_schema, + recon_k, + emitted_row_ids, + scopes, + ) + .await + }) + .try_filter_map(|opt| future::ready(Ok(opt))), + ) .finally(move || { // Publish the exact index-file I/O measured for this query // (cache misses only) to the iops/requests/bytes_read gauges. @@ -2360,13 +2594,24 @@ pub struct MultivectorScoringExec { // the inputs are sorted ANN search results inputs: Vec>, query: Query, + /// Output schema: `[_distance, _rowid]` plus any covering columns the ANN children + /// emit. The covered payload is per source row (replicated across its sub-vectors), + /// so it is carried through the per-row-id reduction below. + output_schema: SchemaRef, properties: Arc, } impl MultivectorScoringExec { pub fn try_new(inputs: Vec>, query: Query) -> Result { + // Inherit the children's schema so covering columns flow through scoring instead + // of being dropped (a covered multivector query would otherwise still take from + // the base table). + let output_schema = inputs + .first() + .map(|input| input.schema()) + .unwrap_or_else(|| KNN_INDEX_SCHEMA.clone()); let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(KNN_INDEX_SCHEMA.clone()), + EquivalenceProperties::new(output_schema.clone()), Partitioning::RoundRobinBatch(1), EmissionType::Final, Boundedness::Bounded, @@ -2375,6 +2620,7 @@ impl MultivectorScoringExec { Ok(Self { inputs, query, + output_schema, properties, }) } @@ -2399,7 +2645,7 @@ impl ExecutionPlan for MultivectorScoringExec { } fn schema(&self) -> arrow_schema::SchemaRef { - KNN_INDEX_SCHEMA.clone() + self.output_schema.clone() } fn children(&self) -> Vec<&Arc> { @@ -2434,11 +2680,18 @@ impl ExecutionPlan for MultivectorScoringExec { .map(|input| input.execute(partition, context.clone())) .collect::>>()?; + // Covering columns (if any) follow `[_distance, _rowid]` in the child schema. + // They are per source row -- identical across a row's sub-vectors -- so we carry + // them through the per-row-id reduction and re-attach one copy per output row. + let output_schema = self.output_schema.clone(); + let n_covered = output_schema.fields().len().saturating_sub(2); + // collect the top k results from each stream, // and max-reduce for each query, // records the minimum distance for each query as estimation. let mut reduced_inputs = stream::select_all(inputs.into_iter().map(|stream| { - stream.map(|batch| { + let output_schema = output_schema.clone(); + stream.map(move |batch| { let batch = batch?; let row_ids = batch[ROW_ID].as_primitive::(); let dists = batch[DIST_COL].as_primitive::(); @@ -2452,9 +2705,15 @@ impl ExecutionPlan for MultivectorScoringExec { .unwrap_or_default(); let mut new_row_ids = Vec::with_capacity(row_ids.len()); let mut new_sims = Vec::with_capacity(row_ids.len()); + let mut keep_indices: Vec = Vec::with_capacity(row_ids.len()); let mut visited_row_ids = HashSet::with_capacity(row_ids.len()); - for (row_id, dist) in row_ids.values().iter().zip(dists.values().iter()) { + for (i, (row_id, dist)) in row_ids + .values() + .iter() + .zip(dists.values().iter()) + .enumerate() + { // the results are sorted by distance, so we can skip if we have seen this row id before if visited_row_ids.contains(row_id) { continue; @@ -2463,14 +2722,20 @@ impl ExecutionPlan for MultivectorScoringExec { new_row_ids.push(*row_id); // it's cosine distance, so we need to convert it to similarity new_sims.push(1.0 - *dist); + keep_indices.push(i as u32); } - let new_row_ids = UInt64Array::from(new_row_ids); - let new_dists = Float32Array::from(new_sims); - let batch = RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![Arc::new(new_dists), Arc::new(new_row_ids)], - )?; + let mut columns: Vec = vec![ + Arc::new(Float32Array::from(new_sims)), + Arc::new(UInt64Array::from(new_row_ids)), + ]; + if n_covered > 0 { + let keep = arrow_array::UInt32Array::from(keep_indices); + for covered_col in batch.columns().iter().skip(2) { + columns.push(arrow::compute::take(covered_col, &keep, None)?); + } + } + let batch = RecordBatch::try_new(output_schema.clone(), columns)?; Ok::<_, DataFusionError>((min_sim, batch)) }) @@ -2482,11 +2747,28 @@ impl ExecutionPlan for MultivectorScoringExec { let stream = stream::once(async move { // at most, we will have k * refine_factor results for each query let mut results = HashMap::with_capacity(k * refactor); + // Covering columns of every reduced batch, kept alive (cheap Arc clones) + // for one zero-copy-per-cell `interleave` gather at emit. The map records + // where each row id's covering values live: (batch position, row position), + // first-seen (the values are per source row, identical across batches). + // Materializing per-cell `ScalarValue`s here instead would copy every + // candidate cell once into the map and again through `iter_to_array` -- + // several MB per query for wide payloads. + let mut covered_batches: Vec> = Vec::new(); + let mut payload_locs: HashMap = HashMap::new(); let mut missed_sim_sum = 0.0; while let Some((min_sim, batch)) = reduced_inputs.try_next().await? { let row_ids = batch[ROW_ID].as_primitive::(); let sims = batch[DIST_COL].as_primitive::(); + if n_covered > 0 { + let batch_pos = covered_batches.len(); + for (i, row_id) in row_ids.values().iter().enumerate() { + payload_locs.entry(*row_id).or_insert((batch_pos, i)); + } + covered_batches.push(batch.columns().iter().skip(2).cloned().collect()); + } + let query_results = row_ids .values() .iter() @@ -2512,18 +2794,47 @@ impl ExecutionPlan for MultivectorScoringExec { missed_sim_sum += min_sim; } - let (row_ids, sims): (Vec<_>, Vec<_>) = results.into_iter().unzip(); + let (row_ids, sims): (Vec, Vec) = results.into_iter().unzip(); let dists = sims - .into_iter() + .iter() // it's similarity, so we need to convert it back to distance .map(|sim| num_queries - sim) .collect::>(); - let row_ids = UInt64Array::from(row_ids); - let dists = Float32Array::from(dists); - let batch = RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![Arc::new(dists), Arc::new(row_ids)], - )?; + let mut columns: Vec = vec![ + Arc::new(Float32Array::from(dists)), + Arc::new(UInt64Array::from(row_ids.clone())), + ]; + if n_covered > 0 && !row_ids.is_empty() { + // Every emitted row id was recorded while scanning the reduced + // batches (both maps are fed from the same iteration), so a miss + // here is a results/payloads desync. Fail loudly rather than + // fall back to nulls: a silent null would masquerade as a real + // covered value and hide the desync from every result-based test. + let locations = row_ids + .iter() + .map(|row_id| { + payload_locs.get(row_id).copied().ok_or_else(|| { + DataFusionError::Internal(format!( + "multivector covering: row id {row_id} missing from \ + covering source" + )) + }) + }) + .collect::>>()?; + for covered_pos in 0..n_covered { + let arrays: Vec<&dyn Array> = covered_batches + .iter() + .map(|cols| cols[covered_pos].as_ref()) + .collect(); + columns.push(arrow::compute::interleave(&arrays, &locations)?); + } + } else { + for covered_pos in 0..n_covered { + let field = output_schema.field(2 + covered_pos); + columns.push(arrow::array::new_empty_array(field.data_type())); + } + } + let batch = RecordBatch::try_new(output_schema.clone(), columns)?; Ok::<_, DataFusionError>(batch) }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -2600,6 +2911,68 @@ mod tests { } } + /// The read path is manifest-authoritative: if an index declares a covering field id + /// the current schema cannot resolve, `ANNIvfSubIndexExec::try_new` must fail loudly + /// rather than silently produce an output schema missing that column -- which would + /// let `TakeExec` wrongly elide the base-table fetch for a column never emitted. + #[tokio::test] + async fn test_ann_sub_index_rejects_unresolvable_covering_field() { + use lance_datafusion::exec::OneShotExec; + + // Minimal input plan carrying the required PART_ID_COLUMN schema. + let input = Arc::new(OneShotExec::from_batch(RecordBatch::new_empty( + KNN_PARTITION_SCHEMA.clone(), + ))); + + // A dataset whose schema has no field with id 9999. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-unresolvable-covering", + None, + ) + .await + .unwrap(), + ); + + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![9999], + }; + + let result = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![index], + base_query(), + PreFilterSource::None, + ); + let err = result.expect_err("unresolvable covering field id must be rejected"); + assert!( + err.to_string().contains("covering field id 9999"), + "expected a manifest/schema inconsistency error, got: {err}" + ); + } + #[test] fn test_effective_query_parallelism_clamps_to_cpu_pool() { let mut query = base_query(); @@ -2826,7 +3199,12 @@ mod tests { } fn find_partitions(&self, _query: &Query) -> Result<(UInt32Array, Float32Array)> { - unimplemented!() + // No additional partitions to reveal; the covered re-rank in + // late_search treats an empty result as "nothing more to search". + Ok(( + UInt32Array::from(Vec::::new()), + Float32Array::from(Vec::::new()), + )) } fn total_partitions(&self) -> usize { @@ -3194,6 +3572,69 @@ mod tests { ); } + /// A `FilterLoader` that yields a fixed allow-list of row addresses, so the + /// resulting prefilter mask is a bounded `AllowList` (`max_len()` is `Some`). + struct StaticAllowList(Vec); + + #[async_trait] + impl FilterLoader for StaticAllowList { + async fn load(self: Box) -> lance_core::Result { + Ok(lance_select::RowAddrMask::from_allowed( + lance_select::RowAddrTreeMap::from_iter(self.0.iter().copied()), + )) + } + } + + /// Like `empty_prefilter`, but the prefilter mask is a bounded allow-list of + /// the given row addresses, which drives the late-search no-rows shortcut. + async fn bounded_prefilter(allow: Vec) -> Arc { + static NEXT_PREFILTER_DATASET_ID: AtomicUsize = AtomicUsize::new(100_000); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let uri = format!( + "memory://bounded-prefilter-{}", + NEXT_PREFILTER_DATASET_ID.fetch_add(1, Ordering::Relaxed) + ); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema), + &uri, + None, + ) + .await + .unwrap(), + ); + let mut indexed_fragments = RoaringBitmap::new(); + for fragment in dataset.manifest.fragments.iter() { + indexed_fragments.insert(fragment.id as u32); + } + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![], + name: "test".to_string(), + dataset_version: 1, + fragment_bitmap: Some(indexed_fragments), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: Vec::new(), + }; + let loader: Box = Box::new(StaticAllowList(allow)); + let prefilter = Arc::new(DatasetPreFilter::new(dataset, &[index], Some(loader))); + prefilter.wait_for_ready().await.unwrap(); + prefilter + } + fn prepared_metrics() -> Arc { Arc::new(AnnIndexMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) } @@ -3407,6 +3848,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>() .await @@ -3522,6 +3964,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3601,6 +4044,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3614,6 +4058,211 @@ mod tests { ); } + /// With a bounded selective prefilter, the late-search no-rows shortcut must fire + /// for covered indexes too -- but emit nothing: the shortcut's `[_distance, _rowid]` + /// batch would drop the covering columns, so the end-of-stream covered recovery + /// (`covered_not_found_batch`) emits the missing rows instead, with a take bounded + /// by the prefilter size. Searching instead of shortcutting is the failure mode this + /// pins down: a prefilter-matched row with no index entry (null vector) can never be + /// found, so the search degenerates to scoring every partition on every query. + #[tokio::test] + async fn test_late_search_covered_skips_search_on_selective_prefilter() { + // Bounded prefilter (one allowed row) + k=2 => max_results=1 <= k, and an + // empty initial state => found_so_far=0 < max_results: shortcut criteria. + async fn run_late_search(has_covered: bool) -> (Vec, Vec) { + let (index, _prepared, searched_partitions, _threads) = + prepared_index(vec![21, 22, 23]); + let mut query = base_query(); + query.k = 2; + query.minimum_nprobes = 0; + query.maximum_nprobes = Some(3); + let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + let prefilter = bounded_prefilter(vec![0]).await; + + let batches = ANNIvfSubIndexExec::late_search( + index, + query, + Arc::new(UInt32Array::from(vec![0, 1, 2])), + Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])), + prefilter.clone(), + prefilter, + prepared_metrics(), + state, + usize::MAX, + None, + has_covered, + ) + .try_collect::>() + .await + .unwrap(); + let searched = searched_partitions.lock().unwrap().clone(); + (batches, searched) + } + + // Control: a non-covered index takes the shortcut, so no partition is + // searched and the only batch is the [_distance, _rowid] not-found batch. + let (batches, searched) = run_late_search(false).await; + assert!( + searched.is_empty(), + "non-covered query should take the shortcut and skip the search, searched={searched:?}" + ); + // The non-covered control for the covered/non-covered contrast below: guard + // against a batch count of zero, which would make `all(..)` vacuously true + // regardless of whether the shortcut actually emitted anything. + assert_eq!(batches.len(), 1, "the shortcut emits exactly one batch"); + assert!( + batches.iter().all(|b| b.num_columns() == 2), + "the shortcut emits only [_distance, _rowid]" + ); + + // Covered: the shortcut fires too (no partition searched), but emits nothing -- + // the end-of-stream covered recovery supplies the not-found rows with their + // covering columns. + let (batches, searched) = run_late_search(true).await; + assert!( + searched.is_empty(), + "covered query must take the shortcut and skip the search, searched={searched:?}" + ); + assert!( + batches.is_empty(), + "the covered shortcut emits nothing; recovery happens at end of stream" + ); + } + + /// Arming must record the arming delta's ownership scope, and the recovery must + /// emit only rows an armed scope selects. A sibling delta that skipped its late + /// search without arming never emits its not-found rows on the non-covered path + /// (its `seg_mask` filters them out of the shortcut batch), so recovering them + /// for a covered index would make the covered and non-covered result sets + /// diverge -- e.g. null-vector rows owned by the skipping delta surfacing at + /// INFINITY distance only when the index is covered. + #[tokio::test] + async fn test_covered_recovery_honors_arming_delta_ownership() { + // Delta B owns row 0 only; rows 1 and 2 belong to a sibling delta that will + // NOT arm (it skipped its late search entirely). + let seg_mask_b = Arc::new(RowAddrMask::from_allowed( + lance_select::RowAddrTreeMap::from_iter([0u64]), + )); + + // Delta B takes the covered no-rows shortcut: bounded prefilter of 3 rows, + // k = 4, nothing found so far. + let (index, _prepared, searched_partitions, _threads) = prepared_index(vec![21, 22, 23]); + let mut query = base_query(); + query.k = 4; + query.minimum_nprobes = 0; + query.maximum_nprobes = Some(3); + let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); + let prefilter = bounded_prefilter(vec![0, 1, 2]).await; + + let batches = ANNIvfSubIndexExec::late_search( + index, + query, + Arc::new(UInt32Array::from(vec![0, 1, 2])), + Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3])), + prefilter.clone(), + prefilter.clone(), + prepared_metrics(), + state.clone(), + usize::MAX, + Some(seg_mask_b.clone()), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert!(batches.is_empty(), "covered shortcut emits nothing inline"); + assert!( + searched_partitions.lock().unwrap().is_empty(), + "the shortcut must not search" + ); + + // Arming recorded delta B's scope (and only that one). + let scopes = std::mem::take(&mut *state.covered_recovery_scopes.lock().unwrap()); + assert_eq!(scopes.len(), 1, "exactly one delta armed"); + assert!(scopes[0].is_some(), "the armed scope is delta B's seg_mask"); + + // The recovery over a real dataset must return ONLY delta B's row 0 -- + // rows 1 and 2 are owned by the never-armed sibling. + let payload_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "payload", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + payload_schema.clone(), + vec![Arc::new(Int32Array::from(vec![10, 11, 12]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), payload_schema), + "memory://covered-recovery-ownership", + None, + ) + .await + .unwrap(), + ); + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + fields.push(Arc::new(ArrowField::new("payload", DataType::Int32, false))); + let output_schema = Arc::new(ArrowSchema::new(fields)); + + let emitted = Arc::new(std::sync::Mutex::new(roaring::RoaringTreemap::new())); + let recovered = ANNIvfSubIndexExec::covered_not_found_batch( + true, + prefilter.clone(), + dataset.clone(), + output_schema.clone(), + 4, + emitted.clone(), + scopes, + ) + .await + .unwrap() + .expect("delta B's unemitted row must be recovered"); + assert_eq!( + recovered + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values(), + &[0], + "recovery must emit only rows the arming delta owns" + ); + assert_eq!( + recovered + .column_by_name("payload") + .unwrap() + .as_primitive::() + .values(), + &[10], + "the covering column value comes from the base table" + ); + + // A global (unrestricted) arming scope recovers every unemitted row -- the + // single-delta behaviour is unchanged. + let recovered = ANNIvfSubIndexExec::covered_not_found_batch( + true, + prefilter, + dataset, + output_schema, + 4, + emitted, + vec![None], + ) + .await + .unwrap() + .expect("a global scope recovers all unemitted rows"); + assert_eq!( + recovered + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values(), + &[0, 1, 2], + "a global arming scope keeps the old recover-everything behaviour" + ); + } + #[tokio::test] async fn test_delta_skipping_late_search_releases_sibling() { let prefilter = empty_prefilter().await; @@ -3634,6 +4283,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>(); @@ -3653,6 +4303,7 @@ mod tests { state, usize::MAX, None, + false, ) .try_collect::>(); @@ -4118,6 +4769,84 @@ mod tests { } } + /// Covered multivector scoring must carry each row's covering values through + /// the per-row-id reduction, aligned by row id -- and must FAIL LOUDLY (not + /// emit nulls) if an emitted row id ever has no recorded covering source. + #[tokio::test] + async fn test_multivector_score_carries_covering_values() { + let query = Query { + column: "vector".to_string(), + key: Arc::new(generate_random_array(1)), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::Cosine), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + }; + + let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); + fields.push(Arc::new(ArrowField::new("payload", DataType::Int32, false))); + let covered_schema = Arc::new(ArrowSchema::new(fields)); + + // Three sub-query streams over overlapping row ids 1..=4; the covering value + // is a pure function of the row id (100 + row_id), identical across streams, + // as the contract requires. + let inputs = (0..3u64) + .map(|i| { + let batch = RecordBatch::try_new( + covered_schema.clone(), + vec![ + Arc::new(Float32Array::from(vec![i as f32 + 1.0, i as f32 + 2.0])), + Arc::new(UInt64Array::from(vec![i + 1, i + 2])), + Arc::new(Int32Array::from(vec![ + 100 + (i + 1) as i32, + 100 + (i + 2) as i32, + ])), + ], + ) + .unwrap(); + let input: Arc = Arc::new(TestingExec::new(vec![batch])); + input + }) + .collect::>(); + + let ctx = Arc::new(datafusion::execution::context::TaskContext::default()); + let plan = MultivectorScoringExec::try_new(inputs, query).unwrap(); + let batches = plan + .execute(0, ctx) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let mut seen = 0; + for batch in batches { + assert_eq!(batch.num_columns(), 3, "covering column must flow through"); + let row_ids = batch[ROW_ID].as_primitive::(); + let payloads = batch + .column_by_name("payload") + .unwrap() + .as_primitive::(); + assert_eq!(payloads.null_count(), 0, "no fabricated nulls"); + for (row_id, payload) in row_ids.values().iter().zip(payloads.values().iter()) { + assert_eq!( + *payload, + 100 + *row_id as i32, + "covering value must stay aligned with its row id" + ); + seen += 1; + } + } + assert_eq!(seen, 4, "all four distinct row ids are emitted"); + } + /// A test dataset for testing the nprobes parameter. /// /// The dataset has 100 partitions and filterable columns setup to easily create From de7f0e5ff45b2e17324c7d04e380e56ea7fb69ed Mon Sep 17 00:00:00 2001 From: Vivek Date: Sun, 16 Aug 2026 20:18:24 -0700 Subject: [PATCH 2/4] feat(index): extend covering columns to every IVF vector index type Covering columns worked only on IVF_PQ. Each storage format now declares which of its own columns are internal, so anything else is treated as covered payload, and the creation-time restriction to IVF_PQ is gone. IVF_SQ, RQ, FLAT and the HNSW variants each get end-to-end coverage. --- rust/lance-index/src/vector/bq/storage.rs | 58 + rust/lance-index/src/vector/flat/storage.rs | 94 +- rust/lance-index/src/vector/quantizer.rs | 10 +- rust/lance-index/src/vector/sq/storage.rs | 64 +- rust/lance-index/src/vector/storage.rs | 16 + .../src/system_index/frag_reuse.rs | 48 +- rust/lance/src/index/vector.rs | 295 ++-- rust/lance/src/index/vector/builder.rs | 38 +- rust/lance/src/index/vector/ivf/v2.rs | 1278 +++++++++++++++-- 9 files changed, 1647 insertions(+), 254 deletions(-) diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 0f7f034b2c1..cf719789e94 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -44,6 +44,7 @@ use crate::frag_reuse::{FragReuseIndex, FragReuseIndexHandle}; use crate::pb; use crate::scalar::RowIdRemapper; use crate::vector::ApproxMode; +use crate::vector::PART_ID_COLUMN; use crate::vector::bq::dist_table_quant::{ DistTableDequant, quantize_dist_table_into, quantize_dist_table_u16_into, }; @@ -66,6 +67,7 @@ use crate::vector::pq::storage::transpose; use crate::vector::quantizer::{QuantizerMetadata, QuantizerStorage}; use crate::vector::storage::{ DistCalculator, DistanceCalculatorOptions, QueryResidual, RabitRawQueryContext, VectorStore, + covering_field_indices_excluding, }; pub const RABIT_METADATA_KEY: &str = "lance:rabit"; @@ -2048,6 +2050,29 @@ impl VectorStore for RabitQuantizationStorage { self.batch.schema_ref() } + /// RQ storage carries the row id plus many internal columns (the binary codes, the + /// extended-bit codes, and the per-row factor columns, some present only for higher + /// `num_bits`); the covering ("included") columns are everything else. + fn covering_field_indices(&self) -> Vec { + // Every RaBitQ-internal column plus the legacy `__ivf_part_id` column + // (left in pre-#3606 single-partition builds; unlike PQ, RaBitQ storage + // does not drop it at construction, so it is excluded here). Anything + // else is a user-declared covering column. + const INTERNAL: &[&str] = &[ + ROW_ID, + RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + PART_ID_COLUMN, + ]; + covering_field_indices_excluding(self.schema().as_ref(), INTERNAL) + } + fn to_batches(&self) -> Result + Send> { Ok(std::iter::once(self.batch.clone())) } @@ -3035,6 +3060,39 @@ mod tests { .unwrap() } + #[test] + fn test_covering_excludes_legacy_part_id_column() { + // Same legacy `__ivf_part_id` guard as the other quantizer storages: + // pre-#3606 `num_partitions=1` builds left it in the storage file, and + // it must never be classified as a covering column -- otherwise search + // emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks. + use crate::vector::PART_ID_COLUMN; + use arrow_array::UInt32Array; + use arrow_schema::{DataType, Field, Schema}; + + let code_dim = 64; + let codes = make_test_codes(10, code_dim); + let metadata = make_test_metadata(codes.value_length() as usize * 8); + let base = make_test_batch(codes); + + let mut fields: Vec<_> = base.schema().fields().iter().cloned().collect(); + fields.push(Arc::new(Field::new(PART_ID_COLUMN, DataType::UInt32, true))); + let mut columns = base.columns().to_vec(); + columns.push(Arc::new(UInt32Array::from_iter_values( + (0..base.num_rows()).map(|_| 0), + )) as ArrayRef); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + + let storage = + RabitQuantizationStorage::try_from_batch(batch, &metadata, DistanceType::L2, None) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + fn make_test_ex_codes(num_vectors: usize, code_dim: usize, num_bits: u8) -> FixedSizeListArray { let ex_bits = rabit_ex_bits(num_bits).unwrap(); let ex_code_bytes = rabit_ex_code_bytes(code_dim, ex_bits).unwrap(); diff --git a/rust/lance-index/src/vector/flat/storage.rs b/rust/lance-index/src/vector/flat/storage.rs index 1de4c4387fb..3c90f0c9e01 100644 --- a/rust/lance-index/src/vector/flat/storage.rs +++ b/rust/lance-index/src/vector/flat/storage.rs @@ -6,8 +6,11 @@ use std::{borrow::Cow, sync::Arc}; use super::index::FlatMetadata; use crate::frag_reuse::{FragReuseIndex, FragReuseIndexHandle}; use crate::scalar::RowIdRemapper; +use crate::vector::PART_ID_COLUMN; use crate::vector::quantizer::QuantizerStorage; -use crate::vector::storage::{DistCalculator, VectorStore}; +use crate::vector::storage::{ + DistCalculator, VectorStore, covering_field_indices_excluding, remap_row_ids_by_name, +}; use crate::vector::utils::do_prefetch; use arrow::array::AsArray; use arrow::compute::concat_batches; @@ -85,7 +88,7 @@ impl QuantizerStorage for FlatFloatStorage { frag_reuse_index: Option>, ) -> Result { let batch = if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } else { batch }; @@ -194,6 +197,17 @@ impl VectorStore for FlatFloatStorage { self.batch.schema_ref() } + /// Flat storage is `[_rowid, flat, ]`, so the covering columns + /// are every field except the row id, the flat vector column, and the legacy + /// `__ivf_part_id` column (left in pre-#3606 single-partition builds; unlike + /// PQ, flat storage does not drop it at construction, so it is excluded here). + fn covering_field_indices(&self) -> Vec { + covering_field_indices_excluding( + self.schema().as_ref(), + &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN], + ) + } + fn as_any(&self) -> &dyn std::any::Any { self } @@ -265,7 +279,7 @@ impl QuantizerStorage for FlatBinStorage { frag_reuse_index: Option>, ) -> Result { let batch = if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } else { batch }; @@ -369,6 +383,17 @@ impl VectorStore for FlatBinStorage { self.batch.schema_ref() } + /// Flat storage is `[_rowid, flat, ]`, so the covering columns + /// are every field except the row id, the flat vector column, and the legacy + /// `__ivf_part_id` column (left in pre-#3606 single-partition builds; unlike + /// PQ, flat storage does not drop it at construction, so it is excluded here). + fn covering_field_indices(&self) -> Vec { + covering_field_indices_excluding( + self.schema().as_ref(), + &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN], + ) + } + fn as_any(&self) -> &dyn std::any::Any { self } @@ -685,6 +710,69 @@ mod tests { FlatFloatStorage::new(vectors, DistanceType::L2) } + // Pre-#3606 `num_partitions=1` builds left `__ivf_part_id` in the flat + // storage file. It must never be classified as a covering column -- + // otherwise search emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks, + // desyncing the two and failing every query on that legacy index. + #[test] + fn test_flat_float_covering_excludes_legacy_part_id_column() { + use crate::vector::PART_ID_COLUMN; + const DIM: i32 = 4; + const N: usize = 8; + let row_ids = UInt64Array::from_iter_values(0..N as u64); + let values = + arrow_array::Float32Array::from_iter_values((0..N * DIM as usize).map(|v| v as f32)); + let vectors = FixedSizeListArray::try_new_from_values(values, DIM).unwrap(); + let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0)); + let batch = RecordBatch::try_from_iter_with_nullable(vec![ + (ROW_ID, Arc::new(row_ids) as ArrayRef, true), + (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true), + (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true), + ]) + .unwrap(); + let storage = FlatFloatStorage::try_from_batch( + batch, + &FlatMetadata { dim: DIM as usize }, + DistanceType::L2, + None, + ) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + + #[test] + fn test_flat_bin_covering_excludes_legacy_part_id_column() { + use crate::vector::PART_ID_COLUMN; + const DIM: i32 = 8; + const N: usize = 8; + let row_ids = UInt64Array::from_iter_values(0..N as u64); + let codes = + arrow_array::UInt8Array::from_iter_values((0..N * DIM as usize).map(|v| v as u8)); + let vectors = FixedSizeListArray::try_new_from_values(codes, DIM).unwrap(); + let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0)); + let batch = RecordBatch::try_from_iter_with_nullable(vec![ + (ROW_ID, Arc::new(row_ids) as ArrayRef, true), + (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true), + (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true), + ]) + .unwrap(); + let storage = FlatBinStorage::try_from_batch( + batch, + &FlatMetadata { dim: DIM as usize }, + DistanceType::L2, + None, + ) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + #[test] fn test_flat_float_storage_distance_f16() { let storage = make_f16_storage(); diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index 433fa5031d4..c8b43636e24 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -290,7 +290,15 @@ pub trait QuantizerStorage: Clone + Sized + DeepSizeOf + VectorStore { let mut indices = Vec::with_capacity(b.num_rows()); let mut new_row_ids = Vec::with_capacity(b.num_rows()); - let row_ids = b.column(0).as_primitive::().values(); + // Look up by name, not position: a covering storage batch is + // `[_rowid, code, ]` for some quantizers but the + // covering columns can precede `_rowid` if the writer schema was + // inferred from a batch rather than declared, so position 0 is + // not reliably the row id column. + let row_id_col = b + .column_by_name(ROW_ID) + .ok_or_else(|| Error::index(format!("column {ROW_ID} not found in batch")))?; + let row_ids = row_id_col.as_primitive::().values(); for (i, row_id) in row_ids.iter().enumerate() { match mapping.get(*row_id) { Some(Some(new_id)) => { diff --git a/rust/lance-index/src/vector/sq/storage.rs b/rust/lance-index/src/vector/sq/storage.rs index 6062c22242e..c7b42c4c33c 100644 --- a/rust/lance-index/src/vector/sq/storage.rs +++ b/rust/lance-index/src/vector/sq/storage.rs @@ -30,9 +30,12 @@ use crate::scalar::RowIdRemapper; use crate::{ INDEX_METADATA_SCHEMA_KEY, IndexMetadata, vector::{ - SQ_CODE_COLUMN, + PART_ID_COLUMN, SQ_CODE_COLUMN, quantizer::{QuantizerMetadata, QuantizerStorage}, - storage::{DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore}, + storage::{ + DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore, + covering_field_indices_excluding, remap_row_ids_by_name, + }, transform::Transformer, }, }; @@ -193,7 +196,7 @@ impl ScalarQuantizationStorage { offsets.push(0); for mut batch in batches.into_iter() { if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { - batch = frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? + batch = remap_row_ids_by_name(batch, frag_reuse_index_ref.as_ref())? } offsets.push(offsets.last().unwrap() + batch.num_rows() as u32); let chunk = SQStorageChunk::new(batch)?; @@ -370,6 +373,17 @@ impl VectorStore for ScalarQuantizationStorage { self.chunks[0].schema() } + /// SQ storage is `[_rowid, __sq_code, ]`, so the covering + /// columns are every field except the row id, the SQ code column, and the + /// legacy `__ivf_part_id` column (left in pre-#3606 single-partition builds; + /// unlike PQ, SQ does not drop it at construction, so it is excluded here). + fn covering_field_indices(&self) -> Vec { + covering_field_indices_excluding( + self.schema().as_ref(), + &[ROW_ID, SQ_CODE_COLUMN, PART_ID_COLUMN], + ) + } + fn as_any(&self) -> &dyn std::any::Any { self } @@ -785,6 +799,50 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(row_ids), Arc::new(code_arr)]).unwrap() } + #[test] + fn test_covering_excludes_legacy_part_id_column() { + // Pre-#3606 `num_partitions=1` builds left `__ivf_part_id` in the SQ + // storage file. It must never be classified as a covering column -- + // otherwise search emits a column the exec's declared schema (from + // `IndexMetadata.covering_fields`, empty for those indexes) lacks, + // desyncing the two and failing every query on that legacy index. + use crate::vector::PART_ID_COLUMN; + use arrow_array::UInt32Array; + const DIM: usize = 64; + const N: usize = 10; + + let row_ids = UInt64Array::from_iter_values(0..N as u64); + let sq_code = UInt8Array::from_iter_values((0..N * DIM).map(|v| v as u8)); + let code_arr = FixedSizeListArray::try_new_from_values(sq_code, DIM as i32).unwrap(); + let part_ids = UInt32Array::from_iter_values((0..N).map(|_| 0)); + + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, false), + Field::new( + SQ_CODE_COLUMN, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt8, true)), + DIM as i32, + ), + false, + ), + Field::new(PART_ID_COLUMN, DataType::UInt32, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(row_ids), Arc::new(code_arr), Arc::new(part_ids)], + ) + .unwrap(); + + let storage = + ScalarQuantizationStorage::try_new(8, DistanceType::L2, -0.7..0.7, [batch], None) + .unwrap(); + assert!( + storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as a covering column" + ); + } + #[test] fn test_get_chunks() { const DIM: usize = 64; diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index fb2be204dd5..9dc843b9d6f 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -392,6 +392,22 @@ pub(crate) fn covering_field_indices_excluding( .collect() } +/// Remap `batch`'s row ids through `frag_reuse_index`, locating the row id column +/// by name rather than assuming it is at a fixed position. A covering storage batch +/// is built from a scan projection of `[, vector]` with `_rowid` +/// appended by the scan, so `_rowid` lands wherever the covering columns end -- +/// never reliably at a fixed index once any are configured. +pub(crate) fn remap_row_ids_by_name( + batch: RecordBatch, + frag_reuse_index: &dyn RowIdRemapper, +) -> Result { + let row_id_idx = batch + .schema() + .index_of(ROW_ID) + .map_err(|_| Error::schema(format!("column {} not found", ROW_ID)))?; + frag_reuse_index.remap_row_ids_record_batch(batch, row_id_idx) +} + /// TODO: should we rename this to "VectorDistance"?; /// ///
diff --git a/rust/lance-table/src/system_index/frag_reuse.rs b/rust/lance-table/src/system_index/frag_reuse.rs index 401cbde7a87..20b0f99fbed 100644 --- a/rust/lance-table/src/system_index/frag_reuse.rs +++ b/rust/lance-table/src/system_index/frag_reuse.rs @@ -397,11 +397,11 @@ impl CompactFragReuseIndex { RoaringTreemap::from_iter(row_ids.iter().filter_map(|addr| self.remap_row_id(addr))) } - /// Remap a record batch that contains a row_id column at index `row_id_idx` - /// Currently this assumes there are only 2 columns in the schema, - /// which is the case for all indexes. - /// For example, for btree, the schema is (value, row_id). - /// For vector index storage, the schema is (row_id, vector). + /// Remap a record batch that contains a row_id column at index `row_id_idx`. + /// Every other column (there may be one, as for scalar indexes -- `(value, + /// row_id)` -- or several, as for a covered vector index's storage -- + /// `(row_id, code, )`) is row-aligned to the surviving, + /// remapped row ids by the same take. pub fn remap_row_ids_record_batch( &self, batch: RecordBatch, @@ -426,8 +426,9 @@ fn remap_row_ids_record_batch( row_id_idx: usize, remap: impl FnOnce(&mut [Option]), ) -> Result { - assert_eq!(batch.schema().fields().len(), 2); - let other_column_idx = 1 - row_id_idx; + // Every column but `_rowid` is carried through row-aligned, however many there are. + // A covered index's storage batch holds its covering ("included") columns next to the + // codes, so this cannot assume the two-column `[_rowid, ]` shape. let row_ids = batch.column(row_id_idx).as_primitive::(); let mut remapped_row_ids = row_ids .values() @@ -436,26 +437,27 @@ fn remap_row_ids_record_batch( .map(Some) .collect::>(); remap(&mut remapped_row_ids); - let (val_indices, new_row_ids): (Vec, Vec) = remapped_row_ids + let (keep_indices, new_row_ids): (Vec, Vec) = remapped_row_ids .iter() .enumerate() .filter_map(|(idx, new_id)| new_id.map(|new_id| (idx as u64, new_id))) .unzip(); - let new_val_indices = UInt64Array::from_iter_values(val_indices); - let new_vals = arrow::compute::take(batch.column(other_column_idx), &new_val_indices, None)?; - - let mut batch_data: Vec<(usize, ArrayRef)> = vec![ - ( - row_id_idx, - Arc::new(UInt64Array::from_iter_values(new_row_ids)) as ArrayRef, - ), - (other_column_idx, Arc::new(new_vals)), - ]; - batch_data.sort_by_key(|(i, _)| *i); - Ok(RecordBatch::try_new( - batch.schema(), - batch_data.into_iter().map(|(_, item)| item).collect(), - )?) + let keep_indices = UInt64Array::from_iter_values(keep_indices); + let new_row_ids: ArrayRef = Arc::new(UInt64Array::from_iter_values(new_row_ids)); + + let columns = batch + .columns() + .iter() + .enumerate() + .map(|(idx, column)| { + if idx == row_id_idx { + Ok(new_row_ids.clone()) + } else { + Ok(arrow::compute::take(column, &keep_indices, None)?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(batch.schema(), columns)?) } fn remap_row_ids_array( diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 8d57d6893fb..ce37c91fe54 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -1138,16 +1138,6 @@ async fn build_vector_index_impl( } } - // Covering is implemented for IVF_PQ in this change; the other IVF vector types land in a - // follow-up. Reject covering on a non-PQ index up front rather than recording - // `covering_fields` the storage cannot serve. - if !params.covering_columns.is_empty() && !matches!(index_type, IndexType::IvfPq) { - return Err(Error::invalid_input( - "covering_columns (covering columns) are currently only supported for IVF_PQ indexes" - .to_string(), - )); - } - match index_type { IndexType::IvfFlat => match element_type { DataType::Float16 | DataType::Float32 | DataType::Float64 => { @@ -2380,7 +2370,8 @@ mod tests { use crate::index::DatasetIndexExt; use arrow_array::Array; use arrow_array::RecordBatch; - use arrow_array::types::{Float32Type, Int32Type}; + use arrow_array::cast::AsArray; + use arrow_array::types::{Float32Type, Int32Type, UInt64Type}; use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use lance_core::utils::tempfile::TempStrDir; use lance_datagen::{BatchCount, RowCount, array}; @@ -2720,17 +2711,86 @@ mod tests { /// rejected at create time: the build would silently ignore the option while /// `create_index` still records `covering_fields` in the manifest, so the exec would /// declare a covered schema the storage cannot emit and every query on the index would - /// fail. + /// fail. This restriction is about the file format, not the index type -- it applies + /// even though covering now works for every IVF vector index type (see + /// `test_covering_columns_supported_for_all_ivf_index_types` below). + #[tokio::test] + async fn test_covering_columns_rejected_for_legacy_index_file_version() { + let test_dir = TempStrDir::default(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.as_str(), None) + .await + .unwrap(); + + let mut params = VectorIndexParams::with_ivf_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + PQBuildParams::new(8, 8), + ) + .version(IndexFileVersion::Legacy) + .clone(); + params.covering_columns(vec!["id".to_string()]); + let err = dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .expect_err("covering_columns on a legacy/non-V3 IVF_PQ build must be rejected"); + assert!( + err.to_string().contains("covering_columns"), + "error should mention covering_columns, got: {err}" + ); + // Fail fast: no index metadata must have been committed. + assert!( + dataset.load_indices().await.unwrap().is_empty(), + "no index should have been created" + ); + } + + /// Covering ("included") columns now work on every IVF vector index type, not just + /// IVF_PQ: each storage format declares which of its own columns are internal (see + /// `VectorStore::covering_field_indices`), so the creation-time restriction that used to + /// reject non-PQ types here (`"covering_columns ... only supported for IVF_PQ"`) is gone. /// - /// Covering is implemented for IVF_PQ only in this change (see the covering validation - /// in `build_vector_index_impl`); every other IVF vector type must reject it too, or a - /// covered index of that type commits `covering_fields` its storage cannot serve and - /// every covered query on it fails at read time instead of at creation time. These - /// non-PQ cases are exactly the ones a would-be `test_initialize_vector_index_preserves_covering` - /// parametrization would need once covering extends beyond IVF_PQ -- until then, they - /// belong here, asserting rejection, not there asserting success. + /// This is the per-type coverage of the *creation* path: `create_index` succeeds and the + /// committed `IndexMetadata.covering_fields` records exactly the requested column. The + /// read-path proof (search reads the covering column back from the index, no base-table + /// take) is covered per-type by `ivf::v2::tests::test_covered_projection_skips_take`. + #[rstest::rstest] + #[case::ivf_pq(VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, 2))] + #[case::ivf_sq(VectorIndexParams::with_ivf_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + SQBuildParams::default(), + ))] + #[case::ivf_hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(8, 8), + ))] + #[case::ivf_hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default(), + ))] + #[case::ivf_rq(VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ))] + #[case::ivf_flat(VectorIndexParams::ivf_flat(4, MetricType::L2))] + #[case::ivf_hnsw_flat(VectorIndexParams::ivf_hnsw( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + ))] #[tokio::test] - async fn test_covering_columns_rejected_for_unsupported_index_types() { + async fn test_covering_columns_supported_for_all_ivf_index_types( + #[case] mut params: VectorIndexParams, + ) { let test_dir = TempStrDir::default(); let reader = lance_datagen::gen_batch() .col("id", array::step::()) @@ -2740,85 +2800,61 @@ mod tests { .await .unwrap(); - const ONLY_IVF_PQ: &str = "only supported for IVF_PQ"; - for (label, mut params, expected_fragment) in [ - ( - "IVF_PQ legacy", - VectorIndexParams::with_ivf_pq_params( - MetricType::L2, - IvfBuildParams::new(4), - PQBuildParams::new(8, 8), - ) - .version(IndexFileVersion::Legacy) - .clone(), - "covering_columns", - ), - ( - "IVF_SQ", - VectorIndexParams::with_ivf_sq_params( - MetricType::L2, - IvfBuildParams::new(4), - SQBuildParams::default(), - ), - ONLY_IVF_PQ, - ), - ( - "IVF_HNSW_PQ", - VectorIndexParams::with_ivf_hnsw_pq_params( - MetricType::L2, - IvfBuildParams::new(4), - HnswBuildParams::default(), - PQBuildParams::new(8, 8), - ), - ONLY_IVF_PQ, - ), - ( - "IVF_HNSW_SQ", - VectorIndexParams::with_ivf_hnsw_sq_params( - MetricType::L2, - IvfBuildParams::new(4), - HnswBuildParams::default(), - SQBuildParams::default(), - ), - ONLY_IVF_PQ, - ), - ( - "IVF_RQ", - VectorIndexParams::with_ivf_rq_params( - MetricType::L2, - IvfBuildParams::new(4), - RQBuildParams::with_rotation_type(1, RQRotationType::Fast), - ), - ONLY_IVF_PQ, - ), - ( - "IVF_FLAT", - VectorIndexParams::ivf_flat(4, MetricType::L2), - ONLY_IVF_PQ, - ), - ( - "IVF_HNSW_FLAT", - VectorIndexParams::ivf_hnsw( - MetricType::L2, - IvfBuildParams::new(4), - HnswBuildParams::default(), - ), - ONLY_IVF_PQ, - ), - ] { - params.covering_columns(vec!["id".to_string()]); - let err = dataset - .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) - .await - .expect_err(&format!("covering_columns on {label} must be rejected")); - assert!( - err.to_string().contains(expected_fragment), - "{label}: error should contain {expected_fragment:?}, got: {err}" - ); - // Fail fast: no index metadata must have been committed. - assert!( - dataset.load_indices().await.unwrap().is_empty(), - "{label}: no index should have been created" + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!( + indices.len(), + 1, + "exactly one index should have been created" + ); + let id_field_id = dataset.schema().field("id").unwrap().id; + assert_eq!( + indices[0].covering_fields, + vec![id_field_id], + "covering_fields must record exactly the requested covering column" + ); + + // The metadata above is derived purely from the request (`create.rs`) and + // never consults the built storage, so it stays correct even if + // `with_covering_columns` were neutralized into a no-op. Force a read that + // only the storage can satisfy: a covered projection must be answered + // without a base-table take, and the returned values must be the row's + // true covering value (not garbage/misaligned) -- proof the storage + // actually carries the column, not just the metadata. + let q = arrow_array::Float32Array::from(vec![0.0f32; 32]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should be answered from the index storage, \ + not a base-table take; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted by the storage") + .as_primitive::(); + let row_ids = batch + .column_by_name(lance_core::ROW_ID) + .expect("row id column") + .as_primitive::(); + assert_eq!(ids.len(), 10, "should return k=10 rows"); + // Single-fragment, step-id dataset => id == row offset == _rowid, so a + // correctly covered id equals the row id for every returned row. + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must be the row's true value, not a no-op/garbage storage column" ); } } @@ -2828,12 +2864,41 @@ mod tests { /// `covering_fields` metadata -- otherwise the target advertises covering /// columns its storage can't emit and covered queries fail. /// - /// IVF_PQ only: covering is implemented only for IVF_PQ in this change (see - /// `build_vector_index_impl`'s covering validation), so this only parametrizes - /// over PQ. Other IVF vector types land coverage in a follow-up that extends - /// `covering_columns` support to them. + /// Covering works on every IVF vector index type (see + /// `test_covering_columns_supported_for_all_ivf_index_types`), and + /// `initialize_vector_index`/`build_vector_index_incremental` thread + /// `covering_columns` through every `(SubIndexType, QuantizationType)` branch with no + /// per-type gate, so this parametrizes over all of them the same way. #[rstest::rstest] #[case::pq(VectorIndexParams::ivf_pq(10, 8, 16, MetricType::L2, 50))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + MetricType::L2, + IvfBuildParams::new(10), + SQBuildParams::default(), + ))] + #[case::rq(VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(10), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast), + ))] + #[case::flat(VectorIndexParams::ivf_flat(10, MetricType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + ))] + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + PQBuildParams::new(8, 8), + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + MetricType::L2, + IvfBuildParams::new(10), + HnswBuildParams::default(), + SQBuildParams::default(), + ))] #[tokio::test] async fn test_initialize_vector_index_preserves_covering( #[case] mut params: VectorIndexParams, @@ -2908,6 +2973,7 @@ mod tests { let query = arrow_array::Float32Array::from(vec![0.5f32; 32]); let mut scan = target_dataset.scan(); scan.nearest("vector", &query, 10).unwrap(); + scan.with_row_id(); scan.project(&["id"]).unwrap(); let plan = scan.explain_plan(true).await.unwrap(); assert!( @@ -2915,10 +2981,25 @@ mod tests { "covered projection should skip the take on the copied index; plan:\n{plan}" ); let batch = scan.try_into_batch().await.unwrap(); - assert!( - batch.column_by_name("id").is_some(), - "copied index should emit the covering column 'id'" - ); + let ids = batch + .column_by_name("id") + .expect("copied index should emit the covering column 'id'") + .as_primitive::(); + let row_ids = batch + .column_by_name(lance_core::ROW_ID) + .expect("row id column") + .as_primitive::(); + // Single-fragment, step-id target dataset => id == row offset == _rowid, so + // a correctly-rebuilt covering column matches the row's true value (not a + // stale copy of the source's storage, and not garbage from a schema + // mismatch during the rebuild). + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "copied index's covered id must be the row's true value, row {i}" + ); + } } #[tokio::test] diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 3121574732d..62db4a77e31 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -1702,6 +1702,12 @@ impl IvfIndexBuilder let quantization_type = Q::quantization_type(); let is_pq = quantization_type == QuantizationType::Product; let is_rq = quantization_type == QuantizationType::Rabit; + // Flat's "code" column is an identity copy of the raw vector (see + // `FlatQuantizer::quantize`), so its true arrow value type (Float16/32/64) + // is only known from the data itself: `FlatQuantizer::field()` hardcodes + // Float32, which is wrong whenever the vector column isn't Float32. Every + // other quantizer type's `field()` is driven by its own parameters (e.g. + // SQ's `num_bits`), not the raw vector's type, and is always correct. let is_flat = quantization_type == QuantizationType::Flat; // prepare the final writers @@ -1710,15 +1716,17 @@ impl IvfIndexBuilder let writer_options = FileWriterOptions::default(); // Covering columns are appended to whichever storage schema this build writes - // (code storage below, or the empty-flat fallback further down). + // (code storage below, or the empty-flat fallback further down). The schema is + // always declared as `[_rowid, code, extra…, covering…]` -- row id first, never + // inferred from a batch's own column order, which can put covering columns + // before `_rowid` and make every downstream reader that trusts column position + // (e.g. the default `QuantizerStorage::remap`) misread row ids. let covering_fields = self.covering_arrow_fields()?; let mut storage_writer = if is_flat { None } else { let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; fields.extend(quantizer.extra_fields()); - // Append any included ("covering") columns so they are persisted in - // the auxiliary storage file alongside the row id and code. fields.extend(covering_fields.iter().cloned()); let storage_schema: Schema = (&arrow_schema::Schema::new(fields)).try_into()?; Some(file_versions::create_writer( @@ -1808,7 +1816,26 @@ impl IvfIndexBuilder } if storage_writer.is_none() { - let storage_schema: Schema = batch.schema_ref().as_ref().try_into()?; + // Only flat reaches here (every other type declared its schema + // eagerly above). Declare `[_rowid, flat, covering…]` explicitly -- + // row id always first -- rather than cloning the batch's own + // schema, whose column order can put covering columns before + // `_rowid`. The flat column's arrow type is looked up by name + // from this batch (not assumed), since it is an identity copy + // of the raw vector column and can be Float16/32/64. + let flat_field = batch + .schema_ref() + .field_with_name(lance_index::vector::flat::storage::FLAT_COLUMN) + .map_err(|e| { + Error::invalid_input(format!( + "flat storage batch missing its code column: {e}" + )) + })? + .clone(); + let mut fields = vec![ROW_ID_FIELD.as_ref().clone(), flat_field]; + fields.extend(covering_fields.iter().cloned()); + let storage_schema: Schema = + (&arrow_schema::Schema::new(fields)).try_into()?; storage_writer = Some(file_versions::create_writer( self.format_version, self.store.create(&storage_path).await?, @@ -1860,6 +1887,9 @@ impl IvfIndexBuilder } if storage_writer.is_none() { + // Every partition was empty, so flat never saw a batch to learn its value + // type from above. The IVF centroids preserve that type (they are computed + // over the same raw vectors `FlatQuantizer::quantize` copies verbatim). let Some(centroids) = ivf.centroids.as_ref() else { return Err(Error::invalid_input( "flat storage writer could not infer schema from empty partitions without IVF centroids", diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 3bd0f34fad5..92c3655827f 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -1128,9 +1128,6 @@ impl IVFIndex { )?) } - /// Append the index's included/covering columns to a per-partition search - /// result (`[_distance, _rowid]`), gathered from the live partition storage. - /// No-op when the index has no covering columns. /// Append the index's included/covering columns to a per-partition search /// result (`[_distance, _rowid]`), gathered from the live partition storage. /// No-op when the index has no covering columns. @@ -3489,143 +3486,715 @@ mod tests { ); } - /// Read-side payoff : a vector query projecting only a covered column - /// is satisfied from the index — no `TakeExec` against the base table. + /// Covering must never change QUERY RESULTS: for any query, a covered index + /// returns exactly the rows and distances a non-covered index returns (covering + /// only changes how projected columns are fetched). Exercises both the late-search + /// shortcut config (min < max nprobes: prefilter-matched rows without index + /// entries come back with INFINITY distance) and the all-partitions-searched + /// corner (min == max: such rows are NOT returned at all). #[tokio::test] - async fn test_ivf_pq_covered_projection_skips_take() { - const INDEX_NAME: &str = "vector_idx"; - let test_dir = TempStrDir::default(); - let test_uri = test_dir.as_str(); - let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + async fn test_covered_prefilter_results_match_non_covered() { + use arrow_array::types::UInt64Type; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; - let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 4096usize; + let n_null = 3usize; // ids 0..3 have NULL vectors (no index entry) + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let ids: Vec = (0..n as i32).collect(); + // The prefilter matches ids 0..5: three null-vector rows plus two real ones. + let cats: Vec<&str> = (0..n) + .map(|i| if i < 5 { "picked" } else { "rest" }) + .collect(); + // Two well-separated clusters (around 0.0 and 100.0). The real picked rows + // (ids 3, 4) live in cluster B; the query below probes cluster A first, so the + // late search / shortcut path is genuinely exercised at min < max nprobes. + let values: Vec = (0..n) + .flat_map(|r| { + let center = if r % 2 == 0 { 0.0f32 } else { 100.0 }; + (0..dim as usize) + .map(move |d| center + ((r * dim as usize + d) % 400) as f32 * 1e-3) + }) + .collect(); + let validity: Vec = (0..n).map(|i| i >= n_null).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_noncovered_parity", + None, + ) + .await + .unwrap(); + + // (rowid -> distance) results for the prefiltered query under the given probe + // configuration and index variant. + async fn run(dataset: &Dataset, min_nprobes: usize, max_nprobes: usize) -> Vec<(u64, f32)> { + let q = Float32Array::from(vec![0.0f32; 4]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.minimum_nprobes(min_nprobes); + scan.maximum_nprobes(max_nprobes); + scan.filter("category = 'picked'").unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the parity comparison requires the index path; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + let row_ids = batch[ROW_ID].as_primitive::(); + let dists = batch["_distance"].as_primitive::(); + let mut out: Vec<(u64, f32)> = row_ids + .values() + .iter() + .zip(dists.values().iter()) + .map(|(r, d)| (*r, *d)) + .collect(); + out.sort_by_key(|(r, _)| *r); + out + } + + for (min_nprobes, max_nprobes) in [(1usize, 2usize), (2, 2)] { + let centroids = || { + let vals: Vec = [0.0f32, 100.0] + .iter() + .flat_map(|c| std::iter::repeat_n(*c, dim as usize)) + .collect(); + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(vals), dim).unwrap(), + ) + }; + let plain_params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(2, centroids()).unwrap(), + ); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + &plain_params, + true, + ) + .await + .unwrap(); + let plain = run(&dataset, min_nprobes, max_nprobes).await; + + let mut covered_params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + IvfBuildParams::try_with_centroids(2, centroids()).unwrap(), + ); + covered_params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + &covered_params, + true, + ) + .await + .unwrap(); + let covered = run(&dataset, min_nprobes, max_nprobes).await; + + assert_eq!( + covered, plain, + "covered results must be identical to non-covered \ + (min_nprobes={min_nprobes}, max_nprobes={max_nprobes})" + ); + } + } + + /// A row whose vector is null has no index entry, so the covered search never returns + /// it -- but a bounded, selective prefilter that admits it still expects it in the + /// results (parity with a non-covered scan). The covered path must recover it, with + /// its covering column fetched from the base table. + #[tokio::test] + async fn test_ivf_covered_recovers_null_vector_prefilter_rows() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 40usize; + let n_rare = 3usize; // category "rare" rows, all with NULL vectors (no index entry) + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let ids: Vec = (0..n as i32).collect(); + let cats: Vec<&str> = (0..n) + .map(|i| if i < n_rare { "rare" } else { "common" }) + .collect(); + // Two well-separated clusters so early pruning keeps minimum_nprobes at 1: + // the recovery path only fires where the non-covered shortcut would (an + // all-partitions-searched query returns only found rows on both). + let values: Vec = (0..n) + .flat_map(|r| { + let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; + (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) + }) + .collect(); + let validity: Vec = (0..n).map(|i| i >= n_rare).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_null_vec_recover", + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); params.covering_columns(vec!["id".to_string()]); dataset .create_index( &["vector"], IndexType::Vector, - Some(INDEX_NAME.to_string()), + Some(INDEX.to_string()), ¶ms, true, ) .await .unwrap(); + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["category"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); - let q = vectors.value(0); - let q = q.as_primitive::(); + // The prefilter admits only the "rare" rows -- all null-vector, so the search + // returns nothing; every result row must come from the recovery path. + let q = Float32Array::from(vec![0.0f32; dim as usize]); let mut scan = dataset.scan(); - scan.nearest("vector", q, 10).unwrap(); - scan.nprobes(4); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter("category = 'rare'").unwrap(); + scan.prefilter(true); scan.with_row_id(); scan.project(&["id"]).unwrap(); - let plan = scan.explain_plan(true).await.unwrap(); - assert!( - !plan.contains("LanceRead"), - "covered projection ['id'] should not require a TakeExec; plan was:\n{plan}" + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + n_rare, + "all null-vector prefilter rows must be recovered (covered parity with non-covered)" ); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for i in 0..ids.len() { + // Row-aligned covered value fetched from the base table for the recovered row. + assert_eq!(ids.value(i) as u64, row_ids.value(i)); + assert!( + (ids.value(i) as usize) < n_rare, + "only the rare rows are admitted" + ); + } + } - let batch = scan.try_into_batch().await.unwrap(); - assert!(batch.column_by_name("id").is_some()); + /// The covered exec emits from two paths in one stream: the search path (rows with an + /// index entry, covering columns read back from index storage) and the recovery path + /// (null-vector rows, covering columns taken from the base table with the exec's declared + /// dataset-typed schema). A prefilter admitting BOTH kinds forces both paths into a single + /// result, so this guards that their batch schemas stay compatible -- if the index + /// round-trip ever changed a covered field's type or nullability, the concatenation here + /// would fail. Covered column is non-nullable to make a nullability drift observable. + #[tokio::test] + async fn test_ivf_covered_mixed_search_and_recovery_share_schema() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; - // Take elision is worthless if the covered search returns the wrong - // neighbors, so also gate on recall: compare the covered result's row ids - // against brute-force ground truth (use_index(false)). All 4 partitions - // are probed, so 0.5 sits far below IVF_PQ's real recall on this data. - let returned: HashSet = batch[ROW_ID] - .as_primitive::() - .values() - .iter() - .copied() + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 40usize; + let n_null = 3usize; // rows 0..3 have NULL vectors (no index entry -> recovery path) + let admit_below = 5i32; // prefilter admits ids 0..5: null rows 0,1,2 + indexed rows 3,4 + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let ids: Vec = (0..n as i32).collect(); + let cats: Vec<&str> = (0..n).map(|_| "x").collect(); + // Two well-separated clusters so early pruning keeps minimum_nprobes at 1 + // (see test_ivf_covered_recovers_null_vector_prefilter_rows). + let values: Vec = (0..n) + .flat_map(|r| { + let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; + (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) + }) .collect(); - let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; - let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; - assert!( - recall >= 0.5, - "covered IVF_PQ recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + let validity: Vec = (0..n).map(|i| i >= n_null).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), ); - } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vector), + ], + ) + .unwrap(); - /// The most likely real query shape: a projection mixing a covered column - /// with an uncovered one. `filtered_read.rs` re-subtracts the projection - /// against the covered stream's schema, so `id` must not be re-fetched - /// while `extra` -- a column the index does not carry -- still goes - /// through a `TakeExec`. Every other covered test projects only-covered - /// or only-uncovered columns and would not catch a regression here. - #[tokio::test] - async fn test_ivf_pq_covered_partial_projection() { - const INDEX_NAME: &str = "vector_idx"; - let test_dir = TempStrDir::default(); - let test_uri = test_dir.as_str(); - let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_mixed_recover", + None, + ) + .await + .unwrap(); - // 16 sub-vectors, not 4: at DIM=32 that is 2 dimensions per codebook instead of - // 8. This fixture is uniform-random, so the true top-10 is full of near-ties, and - // at 4 sub-vectors the PQ error is large enough that NEON-vs-AVX rounding reorders - // them -- the recall gate below read 0.4 on aarch64 while passing on x86. The gate - // is a sanity check on take elision, not a measurement of quantizer accuracy. - let mut params = VectorIndexParams::ivf_pq(4, 8, 16, DistanceType::L2, 2); + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); params.covering_columns(vec!["id".to_string()]); dataset .create_index( &["vector"], IndexType::Vector, - Some(INDEX_NAME.to_string()), + Some(INDEX.to_string()), ¶ms, true, ) .await .unwrap(); - - // A column the index does not cover; the take path must still fetch it. + // A btree on id makes `id < admit_below` a bounded, selective prefilter. + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![( - "extra".to_string(), - "CAST(id AS BIGINT) + 1000".to_string(), - )]), - None, - None, - ) + .create_index(&["id"], IndexType::BTree, None, &scalar, false) .await .unwrap(); - let q = vectors.value(0); - let q = q.as_primitive::(); + let q = Float32Array::from(vec![0.0f32; dim as usize]); let mut scan = dataset.scan(); - scan.nearest("vector", q, 10).unwrap(); - scan.nprobes(4); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter(&format!("id < {admit_below}")).unwrap(); + scan.prefilter(true); scan.with_row_id(); - scan.project(&["id", "extra"]).unwrap(); + scan.project(&["id"]).unwrap(); - let plan = scan.explain_plan(true).await.unwrap(); - assert!( - plan.contains("LanceRead"), - "uncovered column 'extra' must still require a TakeExec; plan was:\n{plan}" + // Concatenating the search-path batches (ids 3,4) with the recovery batch (ids 0,1,2) + // succeeds only if both paths carry a compatible schema for the covered `id` column. + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + admit_below as usize, + "both indexed and null-vector admitted rows must be returned" ); - assert!( - plan.contains("projection=[extra]"), - "the take must fetch only the uncovered column, not re-fetch the \ - already-covered 'id'; plan was:\n{plan}" + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let mut got: Vec = Vec::with_capacity(ids.len()); + for i in 0..ids.len() { + // id == row offset == _rowid, so a row-aligned covered value equals its row id. + assert_eq!(ids.value(i) as u64, row_ids.value(i)); + got.push(ids.value(i)); + } + got.sort_unstable(); + assert_eq!( + got, + (0..admit_below).collect::>(), + "search-path and recovery-path rows must together cover every admitted id" ); + } - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); - let ids = batch["id"].as_primitive::(); - let extras = batch["extra"].as_primitive::(); - for (id, extra) in ids.values().iter().zip(extras.values().iter()) { - assert_eq!( - *extra, - *id as i64 + 1000, - "uncovered 'extra' must match the base table, not a stale/misaligned value" - ); - } + /// On a STABLE-ROW-ID dataset the covered null-vector recovery must resolve stable row + /// ids through the row-id index before taking covering payload -- feeding them to an + /// address-space take (`frag = id >> 32`, `offset = id`) reads the wrong physical row or + /// errors once stable id != physical address. A second fragment makes the two diverge: + /// its rows have addresses `(1 << 32) | offset` but small monotonic stable ids. + #[tokio::test] + async fn test_ivf_covered_recovers_null_vector_stable_row_ids() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator}; + use arrow_buffer::NullBuffer; - // Take elision only matters if the covered result is correct. - let returned: HashSet = batch[ROW_ID] - .as_primitive::() - .values() - .iter() + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let frag0 = 10usize; // fragment 0: ids 0..10, all indexed + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let make_batch = |ids: Vec, null_rows: &[i32]| { + let n = ids.len(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32 + 1.0).collect(); + let validity: Vec = ids.iter().map(|id| !null_rows.contains(id)).collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(validity)), + ); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(vector)], + ) + .unwrap() + }; + + // Fragment 0: ids 0..10, all with vectors. + let batch0 = make_batch((0..frag0 as i32).collect(), &[]); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch0)], schema.clone()), + "memory://covered_stable_null_recover", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: frag0, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Fragment 1 (append): ids 10..14; 10 and 11 have NULL vectors (no index entry). + // Their stable ids (10, 11) are far smaller than their physical addresses + // ((1 << 32) | 0/1), so an address-space take of id 10 lands at fragment 0 offset 10, + // which is out of range (fragment 0 has offsets 0..9) -- the bug. + let batch1 = make_batch((frag0 as i32..frag0 as i32 + 4).collect(), &[10, 11]); + dataset + .append(RecordBatchIterator::new([Ok(batch1)], schema.clone()), None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + // A btree on id makes `id >= 10` a bounded, selective prefilter. + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + // Admit ids 10..14: 12,13 come from the search path, 10,11 from the recovery path. + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + // Recovery matches the non-covered shortcut, which needs unsearched + // partitions to exist (min < max nprobes); all-partitions-searched queries + // return only found rows on both covered and non-covered indexes. + scan.minimum_nprobes(1); + scan.filter("id >= 10").unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4, "all admitted rows must be returned"); + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for i in 0..ids.len() { + // id == stable row id == _rowid, so a correctly-resolved covered value matches. + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must be the row's true value, not an address-space misread" + ); + } + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!(got, vec![10, 11, 12, 13]); + } + + /// Read-side payoff for every vector index type: a query projecting only a covered + /// column is satisfied from the index -- no `TakeExec` against the base table -- + /// with row-aligned values and sane recall. + #[rstest] + #[case::pq(VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + // 16 sub-vectors, not the 4 the plain `pq` case uses: at DIM=32 that is 2 dimensions + // per codebook instead of 8. HNSW *builds* its graph from PQ distances, so a coarse + // quantizer degrades construction and search together, and on this uniform-random + // fixture the true top-10 is full of near-ties -- enough that rounding differences + // between NEON and AVX reorder it, so the recall gate read 0.4 on aarch64 while + // passing on x86. Plain `pq` scans partitions exhaustively and is unaffected, which + // is why only this case needs the extra fidelity. Same reasoning as the `rq` case. + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(16, 8) + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default() + ))] + // RQ uses 5 bits: 1-bit RaBitQ quantization is too coarse to clear the recall + // gate on this random data without a refine (which would re-add the take). + #[case::rq(VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(5, RQRotationType::Fast) + ))] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_projection_skips_take(#[case] mut params: VectorIndexParams) { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + // The HNSW search beam defaults to `k + k / 2` = 15, which on this uniform-random + // fixture leaves the recall gate below a coin flip for the coarsest quantizer: + // `hnsw_pq` returned 0.4 on ~2.5% of runs. The gate is a sanity check on take + // elision, not a measurement of beam width, so give the graph a beam wide enough + // that the assertion tests what it names. Ignored by the non-HNSW cases. + scan.ef(100); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should not require a TakeExec; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert_eq!(ids.len(), 10, "should return k=10 rows"); + // Single-fragment, step-id dataset => id == row offset == _rowid, so a correctly + // covered id equals the row id for every returned row (row-aligned, not stale). + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + row_ids.value(i), + "covered id must match the row's true id (row {i})" + ); + } + + // Take elision is worthless if the covered search returns the wrong neighbors, + // so also gate on recall against brute-force ground truth. All 4 partitions are + // probed, so 0.5 sits far below any quantizer's real recall on this data. + let returned: HashSet = row_ids.values().iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// The most likely real query shape: a projection mixing a covered column + /// with an uncovered one. `filtered_read.rs` re-subtracts the projection + /// against the covered stream's schema, so `id` must not be re-fetched + /// while `extra` -- a column the index does not carry -- still goes + /// through a `TakeExec`. Every other covered test projects only-covered + /// or only-uncovered columns and would not catch a regression here. + #[tokio::test] + async fn test_ivf_pq_covered_partial_projection() { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + // 16 sub-vectors, not 4: at DIM=32 that is 2 dimensions per codebook instead of + // 8. This fixture is uniform-random, so the true top-10 is full of near-ties, and + // at 4 sub-vectors the PQ error is large enough that NEON-vs-AVX rounding reorders + // them -- the recall gate below read 0.4 on aarch64 while passing on x86. The gate + // is a sanity check on take elision, not a measurement of quantizer accuracy. + let mut params = VectorIndexParams::ivf_pq(4, 8, 16, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A column the index does not cover; the take path must still fetch it. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![( + "extra".to_string(), + "CAST(id AS BIGINT) + 1000".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "extra"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "uncovered column 'extra' must still require a TakeExec; plan was:\n{plan}" + ); + assert!( + plan.contains("projection=[extra]"), + "the take must fetch only the uncovered column, not re-fetch the \ + already-covered 'id'; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10, "k=10 should return 10 rows"); + let ids = batch["id"].as_primitive::(); + let extras = batch["extra"].as_primitive::(); + for (id, extra) in ids.values().iter().zip(extras.values().iter()) { + assert_eq!( + *extra, + *id as i64 + 1000, + "uncovered 'extra' must match the base table, not a stale/misaligned value" + ); + } + + // Take elision only matters if the covered result is correct. + let returned: HashSet = batch[ROW_ID] + .as_primitive::() + .values() + .iter() .copied() .collect(); let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; @@ -3715,14 +4284,78 @@ mod tests { /// Same payoff but forcing the BATCH/late-search path: `k` larger than one /// partition makes the initial `minimum_nprobes` sweep under-fill, so /// `late_search` (the run_prepared batch path) fires. Covering must hold there too. + /// + /// Four well-separated clusters, not uniform-random data: on uniform-random + /// data at this `k` (500, i.e. `k > 10`), `early_pruning`'s generous 81x + /// factor (`knn.rs`) makes the pruning heuristic select every partition + /// regardless of the caller's `minimum_nprobes(1)`, so `adjust_probes` bumps + /// `minimum_nprobes` up to `maximum_nprobes` and `late_search` returns empty + /// immediately (`max_nprobes <= min_nprobes`) -- confirmed by instrumenting + /// `late_search` on the prior uniform-random version of this test, which + /// logged `min_nprobes=4 max_nprobes=4` despite the explicit `min=1,max=4` + /// request. With separated clusters the nearest centroid is far closer than + /// the rest, so the heuristic keeps `minimum_nprobes` at 1 and `late_search` + /// genuinely runs. #[tokio::test] async fn test_ivf_pq_covered_projection_batch_path() { const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = NUM_ROWS / NUM_CLUSTERS; + let offsets = [0.0f32, 1000.0, 2000.0, 3000.0]; + + let mut rng = StdRng::seed_from_u64(7); + let mut ids = Vec::with_capacity(NUM_ROWS); + let mut values = Vec::with_capacity(NUM_ROWS * DIM); + for (cluster_idx, offset) in offsets.iter().enumerate() { + for row in 0..ROWS_PER_CLUSTER { + ids.push((cluster_idx * ROWS_PER_CLUSTER + row) as u64); + for dim in 0..DIM { + let base = if dim == 0 { *offset } else { 0.0 }; + let noise = (rng.random::() - 0.5) * 0.02; + values.push(base + noise); + } + } + } + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + false, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM as i32) + .unwrap(), + ), + ], + ) + .unwrap(); + let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); - let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); - let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); + let centroids = build_centroids_for_offsets(&offsets); + let ivf_params = IvfBuildParams::try_with_centroids(NUM_CLUSTERS, centroids).unwrap(); + let mut params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + ivf_params, + PQBuildParams { + num_bits: 8, + num_sub_vectors: 4, + max_iters: 2, + ..Default::default() + }, + ); params.covering_columns(vec!["id".to_string()]); dataset .create_index( @@ -3735,11 +4368,13 @@ mod tests { .await .unwrap(); - let q = vectors.value(0); - let q = q.as_primitive::(); + // A query at cluster 0's exact center: its nearest centroid is a near-zero + // distance away, the rest are ~1000+ away, so `early_pruning` keeps + // `minimum_nprobes` at 1 rather than escalating it to `maximum_nprobes`. + let q = Float32Array::from(vec![0.0f32; DIM]); let mut scan = dataset.scan(); // k >> one partition's rows (NUM_ROWS=512 / 4 parts) forces late expansion. - scan.nearest("vector", q, 500).unwrap(); + scan.nearest("vector", &q, 500).unwrap(); scan.minimum_nprobes(1); scan.maximum_nprobes(4); scan.project(&["id"]).unwrap(); @@ -3751,7 +4386,11 @@ mod tests { ); let batch = scan.try_into_batch().await.unwrap(); assert!(batch.column_by_name("id").is_some()); - assert!(batch.num_rows() > 100, "late search should have expanded"); + assert!( + batch.num_rows() > ROWS_PER_CLUSTER, + "late search should have expanded beyond the nearest partition's {ROWS_PER_CLUSTER} rows, got {}", + batch.num_rows() + ); } /// A covered query combined with a selective prefilter (scalar index + @@ -3830,14 +4469,43 @@ mod tests { /// through the incremental optimize pipeline -- the unindexed-fragment /// shuffle, the partition-split reshuffle, and the partition-join /// reassignment -- as passenger columns (re-gathered by row id). + #[rstest] + #[case::pq(VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[case::hnsw_pq(VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + PQBuildParams::new(4, 8) + ))] + #[case::hnsw_sq(VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + SQBuildParams::default() + ))] + #[case::rq(VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::with_rotation_type(1, RQRotationType::Fast) + ))] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] #[tokio::test] - async fn test_ivf_pq_covered_survives_optimize() { + async fn test_covered_survives_optimize(#[case] mut params: VectorIndexParams) { const INDEX_NAME: &str = "vector_idx"; let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; - let mut params = VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2); params.covering_columns(vec!["id".to_string()]); dataset .create_index( @@ -3857,19 +4525,16 @@ mod tests { .await .unwrap(); - // The merged partition storage must still carry the covering column. - let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; - let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); - assert!( - storage.batch().column_by_name("id").is_some(), - "merged storage should still carry covering column 'id'" - ); - - // And a covered projection is still answered from the index (no take). + // A covered projection must still be answered from the index with no take after + // the merge. This is the real proof that the appended rows kept their covering + // column: if the merge had dropped it, `try_into_batch` would error on the + // schema mismatch (the exec declares `id` from `covering_fields`, but the + // storage could not emit it). let q = vectors.value(0); let q = q.as_primitive::(); let mut scan = dataset.scan(); scan.nearest("vector", q, 10).unwrap(); + scan.with_row_id(); scan.project(&["id"]).unwrap(); let plan = scan.explain_plan(true).await.unwrap(); assert!( @@ -3877,7 +4542,41 @@ mod tests { "covered projection ['id'] should skip the take after optimize; plan:\n{plan}" ); let batch = scan.try_into_batch().await.unwrap(); - assert!(batch.column_by_name("id").is_some()); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + + // Ground truth via an independent (non-index, base-table take) path. + // The dataset now spans two fragments (the original write plus the + // appended delta), so id no longer equals _rowid by construction the + // way it does for the single-fragment `test_covered_projection_skips_take` + // -- but the covered value must still equal the row's *true* id. This + // catches a merge that carries the column but scrambles which row + // its values ended up attached to (e.g. an off-by-one row-id gather). + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id {} != true id {} -- merge scrambled the \ + covering column", + row_ids.value(i), + ids.value(i), + truth_ids.value(i) + ); + } } /// Covering must survive a *retrain* optimize. Retrain rebuilds the storage from @@ -3933,6 +4632,91 @@ mod tests { assert!(batch.column_by_name("id").is_some()); } + /// The streaming partition-search branch (used by HNSW sub-indexes and controlled + /// late searches) must declare the covered schema on the stream it returns, matching + /// the widened batches it emits -- the global-heap branch already does. A stream + /// whose declared schema disagrees with its batches is a latent hazard for any + /// consumer that trusts the declaration. + #[tokio::test] + async fn test_covered_search_partitions_stream_declares_covering_schema() { + use arrow_array::UInt32Array; + use futures::TryStreamExt; + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + let mut params = VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default(), + ); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let query = Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 5, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + }; + let partitions = Arc::new(UInt32Array::from(vec![0u32, 1, 2, 3])); + let dists = Arc::new(Float32Array::from(vec![0.0f32; 4])); + let stream = ctx + .index + .clone() + .search_partitions( + query, + partitions, + dists, + 0, + 4, + Arc::new(NoFilter), + None, + Arc::new(NoOpMetricsCollector), + ) + .await + .unwrap(); + + let declared = stream.schema(); + assert!( + declared.column_with_name("id").is_some(), + "the covered stream must declare the covering column; declared: {declared:?}" + ); + let batches: Vec = stream.try_collect().await.unwrap(); + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!( + batch.schema(), + declared, + "emitted batches must match the declared stream schema" + ); + } + } + /// Index storage that physically carries covering columns while the manifest /// declares none (`covering_fields` empty) must still be queryable: the extra /// columns are dropped from search batches (with a warning) instead of emitting @@ -4188,6 +4972,274 @@ mod tests { } } + /// A covered index with TWO covering columns of different arrow types must + /// survive compaction (which remaps the index's row ids). Every other + /// covering test in this suite uses a single `UInt64` covering column + /// (`id`), which hid two bugs in the default `QuantizerStorage::remap` + /// (used by flat, HNSW_FLAT and SQ, unlike PQ/RQ which override it): + /// it read the row id column by *position* (`column(0)`), which is only + /// ever correct by coincidence for a single leading `UInt64` covering + /// column. With a non-`UInt64` covering column first, the downcast + /// panics; with a `UInt64` one first (as `id` always was), it silently + /// overwrites the storage's real row ids with the covering column's + /// values. Covering column order here is deliberately `[tag, id]` so + /// `tag` (`Int32`) lands first in the storage batch, reproducing the + /// panicking variant, while `id` (`UInt64`) still covers the silent one. + #[rstest] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::hnsw_flat(VectorIndexParams::ivf_hnsw( + DistanceType::L2, + IvfBuildParams::new(4), + HnswBuildParams::default() + ))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_survives_compaction_multiple_covering_columns( + #[case] mut params: VectorIndexParams, + ) { + use arrow_array::Int32Array; + use arrow_array::types::Int32Type; + + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let ids = UInt64Array::from_iter_values(0..NUM_ROWS as u64); + // Distinct from `id` (negative, offset) so a bug that swaps or + // misaligns the two covering columns cannot pass by coincidence. + let tags = Int32Array::from_iter_values((0..NUM_ROWS as i32).map(|i| -i - 1)); + let vectors = generate_random_array_with_range::(NUM_ROWS * DIM, 0.0..1.0); + let fsl = + normalize_fsl(&FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap()) + .unwrap(); + let query_vector = fsl.value(0); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("tag", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(ids), Arc::new(tags), Arc::new(fsl)], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + params.covering_columns(vec!["tag".to_string(), "id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows so compaction rewrites fragments and remaps the index. + dataset.delete("id < 100").await.unwrap(); + compact_after_deletions(&mut dataset).await; + + let q = query_vector.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["tag", "id"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "covered projection ['tag', 'id'] should skip the take after compaction; plan:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + // Plenty of non-deleted rows remain for k=10 to fill; guard against an + // empty result, which would make the checks below vacuously true. + assert_eq!(batch.num_rows(), 10); + + let returned_ids = batch + .column_by_name("id") + .expect("covered projection should return id") + .as_primitive::(); + let returned_tags = batch + .column_by_name("tag") + .expect("covered projection should return tag") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + for v in returned_ids.values() { + assert!(*v >= 100, "deleted rows (id < 100) must not be returned"); + } + + // Ground truth via an independent (non-index, base-table take) path, + // for BOTH covering columns: a corrupted or misaligned remap of + // either column must be caught, not just detected by its presence. + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = + crate::dataset::ProjectionRequest::from_columns(["id", "tag"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + let truth_tags = truth + .column_by_name("tag") + .unwrap() + .as_primitive::(); + for i in 0..returned_ids.len() { + assert_eq!( + returned_ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id after compaction \ + (multi-covering-column case)", + row_ids.value(i) + ); + assert_eq!( + returned_tags.value(i), + truth_tags.value(i), + "row {i} (row_id {}): covered tag != true tag after compaction \ + (multi-covering-column case)", + row_ids.value(i) + ); + } + } + + /// A covered FLAT/SQ index must survive a *deferred* remap: `compact_files` + /// with `defer_index_remap: true` rewrites fragments but leaves the index's + /// row ids to be remapped later via a fragment-reuse index (FRI) instead of + /// remapping them inline. `IvfQuantizationStorage::load_partition` passes + /// that FRI on every subsequent load (search, optimize, another build), so + /// a covered FLAT/SQ storage batch -- `[_rowid, code, ]`, three + /// or more columns -- must not be rejected as though it could only ever be + /// the two-column `(value, row_id)` shape a plain scalar index has. + /// + /// PQ and RQ are excluded: PQ remaps inline via `rebuild_storage_batch` and + /// RQ via its own `remap`, so neither ever reaches the shared FRI path this + /// guards. The HNSW_FLAT/HNSW_SQ variants are also excluded here: HNSW's + /// graph independently fails to survive `defer_index_remap` even without + /// any covering columns (pre-existing, unrelated to covering -- see the + /// P0 fix report), so they would fail this test for a reason this fix + /// does not address. + #[rstest] + #[case::flat(VectorIndexParams::ivf_flat(4, DistanceType::L2))] + #[case::sq(VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + IvfBuildParams::new(4), + SQBuildParams::default() + ))] + #[tokio::test] + async fn test_covered_survives_deferred_remap_frag_reuse( + #[case] mut params: VectorIndexParams, + ) { + const INDEX_NAME: &str = "vector_idx"; + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Delete rows and defer the index remap: this fragment-reuse index now + // covers the covered vector index built above, so every subsequent load + // of its storage carries it. + dataset.delete("id < 100").await.unwrap(); + compact_files( + &mut dataset, + CompactionOptions { + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + // A search must load the (now FRI-tagged) covered storage without + // rejecting its three-or-more-column shape, and still return correct, + // row-aligned covering values for the surviving rows. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + + // Ground truth via an independent (non-index, base-table take) path, + // keyed by the post-compaction row id the scan returned. Compaction + // reassigns physical row addresses, so `id == _rowid` no longer holds + // here the way it does pre-compaction -- but the covered value must + // still equal the row's true id. + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert!( + ids.value(i) >= 100, + "deleted rows (id < 100) must not be returned" + ); + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id after deferred-remap compaction", + row_ids.value(i) + ); + } + + // Building a SECOND covered index while a fragment-reuse index already + // exists on the dataset must also succeed: `StorageBuilder::build` + // threads the FRI into `Q::Storage::try_from_batch` for every new + // build, indexed or not. + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vector_idx_2".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + } + /// Fix 2: under query parallelism > 1 the parallel search branch (which uses /// `search_in_partition`, not the heap merge) must also emit covering /// columns. On a multi-core runner this exercises the parallel path; if the From 7d0ac19d913222b1d5dabf63e56a14df8d3399a9 Mon Sep 17 00:00:00 2001 From: Vivek Date: Wed, 26 Aug 2026 10:45:42 -0700 Subject: [PATCH 3/4] feat(index): keep covering consistent as the data underneath it changes A covered index now survives the operations that change its data: schema evolution, overlays, remap, compaction, and concurrent commits. Guards reject the alterations that would desync covering data -- casting an indexed key column, and any rename, cast or nullability change of a covering column -- because the read path resolves covering columns from the live schema while index storage still emits the old name and type. A partial `merge_insert` that updates a covered column is rewritten as a row-move on the indexed-scan path, so only the rows it touched leave the covering index instead of the whole fragment. The move is an optimisation and never fails the operation: stable row ids, sources carrying inserts, legacy v1 blob columns and partial struct subschemas each fall back to the in-place path, which is correct for all of them. Four public surfaces reported a covered index wrongly -- one answering "no index on this column" for an indexed column, another returning an unrelated column's centroids. All are the same mistake: `fields` answers "what invalidates this index", not "what it can serve". The keyed prefix answers the second. --- .../org/lance/index/IndexDescription.java | 9 +- python/src/indices.rs | 6 +- rust/lance-index/src/traits.rs | 8 +- .../src/format/overlay/staleness.rs | 50 + .../src/transaction/index_maintenance.rs | 452 +++- .../src/transaction/manifest_build.rs | 400 ++- rust/lance/src/dataset/optimize/remapping.rs | 13 +- rust/lance/src/dataset/schema_evolution.rs | 544 +++- rust/lance/src/dataset/take.rs | 13 + rust/lance/src/dataset/write/merge_insert.rs | 2360 ++++++++++++++++- rust/lance/src/index.rs | 564 +++- rust/lance/src/index/vector/ivf/v2.rs | 374 ++- rust/lance/src/io/commit/conflict_resolver.rs | 671 ++++- rust/lance/src/io/exec/knn.rs | 256 +- 14 files changed, 5495 insertions(+), 225 deletions(-) diff --git a/java/src/main/java/org/lance/index/IndexDescription.java b/java/src/main/java/org/lance/index/IndexDescription.java index b1e2689a00c..09a182c192b 100755 --- a/java/src/main/java/org/lance/index/IndexDescription.java +++ b/java/src/main/java/org/lance/index/IndexDescription.java @@ -69,7 +69,14 @@ public String getName() { return name; } - /** Field ids that this index is built on. */ + /** + * Field ids that this index is built on -- the columns it can answer queries for. + * + *

This is the index's keyed prefix only. An index may additionally carry values for + * columns it is not keyed on; those are deliberately absent here, because the index cannot be + * searched on them. They stay reachable per segment via {@link Index#coveringFields()} on the + * entries of {@link #getMetadata()}. + */ public List getFieldIds() { return fieldIds; } diff --git a/python/src/indices.rs b/python/src/indices.rs index acb15491fcb..164dee21ac6 100644 --- a/python/src/indices.rs +++ b/python/src/indices.rs @@ -674,9 +674,11 @@ pub struct PyIndexDescription { pub type_url: String, /// The short type of the index (may not be unique) pub index_type: String, - /// The ids of the fields that the index is built on + /// The ids of the fields the index is keyed on -- the columns it can answer queries about. + /// Covering ("included") columns are excluded; those are reported per-segment as + /// `covering_fields`. pub fields: Vec, - /// The full paths of the fields that the index is built on + /// The full paths of the fields the index is keyed on, matching `fields` /// (dotted, with backtick-quoted segments for non-identifier names) pub field_names: Vec, /// The number of rows indexed by the index diff --git a/rust/lance-index/src/traits.rs b/rust/lance-index/src/traits.rs index f5441bd80f0..9980530d278 100644 --- a/rust/lance-index/src/traits.rs +++ b/rust/lance-index/src/traits.rs @@ -344,7 +344,13 @@ pub trait IndexDescription: Send + Sync { /// deleted. fn rows_indexed(&self) -> u64; - /// Returns the ids of the fields that the index is built on. + /// Returns the ids of the fields that the index is built on -- the columns it can + /// answer queries for. + /// + /// This is the index's *keyed* prefix only. An index may additionally carry values + /// for columns it is not keyed on (see [`IndexMetadata::covering_fields`]); those are + /// deliberately absent here, because the index cannot be searched on them. They stay + /// reachable per segment via [`Self::metadata`]. fn field_ids(&self) -> &[u32]; /// Returns a JSON string representation of the index details diff --git a/rust/lance-table/src/format/overlay/staleness.rs b/rust/lance-table/src/format/overlay/staleness.rs index a6eebe61254..4e096b25ce4 100644 --- a/rust/lance-table/src/format/overlay/staleness.rs +++ b/rust/lance-table/src/format/overlay/staleness.rs @@ -108,6 +108,12 @@ pub fn overlaid_fragments(fragments: &[Fragment]) -> HashMap { /// stale because an overlay committed after the segment was built touches a field the segment /// indexes. Field-aware and version-gated via [`overlay_exclusion_offsets`]. /// +/// Walking `segment.fields` alone covers both the key fields and any covering +/// (`covering_fields`) columns materialized in the index storage: `fields` always lists +/// `covering_fields` as its trailing entries (see [`IndexMetadata::covering_fields`]). An +/// overlay touching either makes the segment's stored copy stale, so both must gate overlay +/// exclusion. +/// /// `overlaid_frags` holds only the fragments that actually carry overlays (rare), so the loop is /// `O(overlaid_frags)` rather than `O(fragments the segment covers)`. pub fn collect_overlay_stale_frags( @@ -412,4 +418,48 @@ mod tests { "fragment absent from bitmap contributes no rows" ); } + + #[test] + fn test_collect_rows_flags_overlay_on_covered_column() { + let schema = flat_test_schema(); + // Segment indexes field 3 (key) and COVERS field 2 (included). An overlay touching + // ONLY the covered field must still flag its rows stale -- otherwise a covered query + // serves the stale index-stored copy of that column instead of the fresh overlay value. + let fragment = fragment_with_overlay(3, dense_overlay(vec![2], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + // `covering_fields` is always the trailing entries of `fields`. + let seg = IndexMetadata { + covering_fields: vec![2], + ..segment(vec![3, 2], 1, Some(bitmap([3]))) + }; + + let mut stale = HashMap::new(); + collect_overlay_stale_rows_for_segment(&seg, &overlaid, &mut stale, &schema).unwrap(); + assert_eq!( + stale.get(&3), + Some(&bitmap([1, 2])), + "overlay on a covered (included) column must flag its rows stale" + ); + } + + #[test] + fn test_collect_frags_flags_overlay_on_covered_column() { + let schema = flat_test_schema(); + // Same at fragment granularity: an overlay on a covered column marks the fragment stale. + let fragment = fragment_with_overlay(3, dense_overlay(vec![2], [1, 2], 9)); + let overlaid: HashMap = HashMap::from([(3u32, &fragment)]); + // `covering_fields` is always the trailing entries of `fields`. + let seg = IndexMetadata { + covering_fields: vec![2], + ..segment(vec![3, 2], 1, Some(bitmap([3]))) + }; + + let mut stale = RoaringBitmap::new(); + collect_overlay_stale_frags(&seg, &overlaid, &mut stale, &schema).unwrap(); + assert_eq!( + stale, + bitmap([3]), + "overlay on a covered (included) column must flag the fragment stale" + ); + } } diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index 06cf9648a88..f96833fe8a3 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -19,6 +19,18 @@ use lance_core::{Error, Result}; use roaring::RoaringBitmap; use std::collections::{HashMap, HashSet}; +/// Insert `field`'s id and every descendant's id into `field_ids`. +/// +/// A private copy of `lance::index::collect_subtree_field_ids`: the walk is four +/// lines and the alternative is exporting a lance-table helper solely so the lance +/// crate's copy could delegate to it. +fn collect_subtree_field_ids(field: &lance_core::datatypes::Field, field_ids: &mut HashSet) { + field_ids.insert(field.id); + for child in &field.children { + collect_subtree_field_ids(child, field_ids); + } +} + impl Transaction { pub(super) fn register_pure_rewrite_rows_update_frags_in_indices( indices: &mut [IndexMetadata], @@ -32,6 +44,13 @@ impl Transaction { return Ok(()); } + // Read the parameter name as the inverse of what it does. Despite + // "for_preserving_frag_bitmap", this list is the set of fields whose VALUES were + // updated by the rewrite, and an index depending on any of them is *excluded* from + // the bitmap extension below -- i.e. passing a field here withholds preservation + // rather than granting it. A caller that cannot prove any field survived the rewrite + // unchanged therefore passes *every* field id, which correctly extends no bitmap at + // all. The name is kept only because renaming it would churn shared callers. let value_updated_field_set = fields_for_preserving_frag_bitmap .iter() .collect::>(); @@ -42,9 +61,21 @@ impl Transaction { if index.results_are_row_addrs() { continue; } - let index_covers_modified_field = index.fields.iter().any(|field_id| { - value_updated_field_set.contains(&u32::try_from(*field_id).unwrap()) - }); + // Covering (`covering_fields`) values are materialized in the index + // storage just like the indexed `fields`, so a rewrite of either + // makes this fragment's entries stale. Reusing them would serve the + // pre-update value, so treat both as "modified field" here. + // + // Expand to the leaf subtree, exactly as every sibling prune does: + // `covering_fields` can only ever hold TOP-LEVEL ids (index creation + // resolves whole columns and rejects dotted names), while the caller + // may report the individual leaves it rewrote. Comparing raw ids would + // miss a covered struct whose child was rewritten and wrongly extend + // this index's bitmap onto the moved fragment, serving the pre-update + // subfield value from index storage. + let index_covers_modified_field = Self::index_dependent_field_ids(index, schema) + .iter() + .any(|field_id| value_updated_field_set.contains(field_id)); if index_covers_modified_field { continue; } @@ -86,27 +117,92 @@ impl Transaction { Ok(()) } - /// If an operation modifies one or more fields in a fragment then we need to remove - /// that fragment from any indices that cover one of the modified fields. + /// The full set of leaf field ids an index depends on: its indexed key fields plus + /// every covered (included) field, each expanded to its whole subtree. `fields` + /// already lists `covering_fields` as its trailing entries (see + /// [`IndexMetadata::covering_fields`]), so walking `fields` alone covers both; a + /// modified/overlaid data file lists leaf ids, so a covered struct's parent id must + /// still be expanded to recognize a change to one of its subfields. Returned as raw + /// `i32` (the field id space of `DataFile.fields` / overlay fields). Shared by the + /// freshness prunes and the conflict-rebase checks. + pub fn index_dependent_leaf_ids(index: &IndexMetadata, schema: &Schema) -> HashSet { + let mut ids: HashSet = HashSet::new(); + for &id in index.fields.iter() { + match schema.field_by_id(id) { + Some(field) => collect_subtree_field_ids(field, &mut ids), + None => { + ids.insert(id); + } + } + } + ids + } + + fn index_dependent_field_ids(index: &IndexMetadata, schema: &Schema) -> HashSet { + Self::index_dependent_leaf_ids(index, schema) + .into_iter() + .filter_map(|id| u32::try_from(id).ok()) + .collect() + } + + /// Drop `updated_fragments` from the coverage of any index that depends on a field + /// listed in `fields_modified`. Covered struct fields are expanded to their leaf + /// subtree: `covering_fields` records the parent struct id while a modified data + /// file lists leaf ids, so exact-id matching would miss an update to a covered + /// struct's subfield and serve its stale value from the index. Every caller must + /// therefore supply the schema (the conflict-rebase path captures it at the read + /// version). pub fn prune_updated_fields_from_indices( indices: &mut [IndexMetadata], updated_fragments: &[Fragment], fields_modified: &[u32], + schema: &Schema, + ) { + if fields_modified.is_empty() { + return; + } + + let deps = Self::index_dependent_field_ids_for_each(indices, schema); + Self::prune_updated_fields_with_deps(indices, updated_fragments, fields_modified, &deps); + } + + /// The dependent-leaf-id set for each index, positionally parallel to `indices`. + /// + /// Resolving these walks the schema and allocates per index, so callers that prune + /// fragment-by-fragment must compute this once and reuse it rather than paying it per + /// fragment -- on a compaction rewriting N fragments across M indices the per-fragment + /// form is N*M schema walks on the commit path. + fn index_dependent_field_ids_for_each( + indices: &[IndexMetadata], + schema: &Schema, + ) -> Vec> { + indices + .iter() + .map(|index| Self::index_dependent_field_ids(index, schema)) + .collect() + } + + /// [`Self::prune_updated_fields_from_indices`] with the per-index dependent-id sets + /// already resolved. `deps` must be positionally parallel to `indices`. + fn prune_updated_fields_with_deps( + indices: &mut [IndexMetadata], + updated_fragments: &[Fragment], + fields_modified: &[u32], + deps: &[HashSet], ) { if fields_modified.is_empty() { return; } + debug_assert_eq!(indices.len(), deps.len()); // If we modified any fields in the fragments then we need to remove those fragments // from the index if the index covers one of those modified fields. - let fields_modified_set = fields_modified.iter().collect::>(); - for index in indices.iter_mut() { - if index - .fields + let fields_modified_set = fields_modified.iter().copied().collect::>(); + for (index, dependent_ids) in indices.iter_mut().zip(deps) { + let touches_index = dependent_ids .iter() - .any(|field_id| fields_modified_set.contains(&u32::try_from(*field_id).unwrap())) - && let Some(fragment_bitmap) = &mut index.fragment_bitmap - { + .any(|field_id| fields_modified_set.contains(field_id)); + if touches_index && let Some(fragment_bitmap) = &mut index.fragment_bitmap { for fragment_id in updated_fragments.iter().map(|f| f.id as u32) { fragment_bitmap.remove(fragment_id); } @@ -116,7 +212,9 @@ impl Transaction { /// Map each (non-tombstoned) field id in a fragment to the path of the data /// file that backs it. - fn fragment_field_paths(frag: &Fragment) -> HashMap { + // `pub`: the conflict-rebase path in the `lance` crate resolves per-fragment + // field ownership with this too. + pub fn fragment_field_paths(frag: &Fragment) -> HashMap { let mut map = HashMap::new(); for file in &frag.files { for &field_id in file.fields.iter() { @@ -138,9 +236,13 @@ impl Transaction { indices: &mut [IndexMetadata], prev_fragments: &[Fragment], new_fragments: &[Fragment], + schema: &Schema, ) { let prev_by_id: HashMap = prev_fragments.iter().map(|f| (f.id, f)).collect(); + // Resolved once for all fragments: this is the commit path, and the per-fragment + // form costs a schema walk and two HashSet allocations per index per fragment. + let deps = Self::index_dependent_field_ids_for_each(indices, schema); for new_frag in new_fragments { let Some(prev) = prev_by_id.get(&new_frag.id) else { continue; // brand-new fragment: nothing stale to prune @@ -160,10 +262,11 @@ impl Transaction { if changed.is_empty() { continue; } - Self::prune_updated_fields_from_indices( + Self::prune_updated_fields_with_deps( indices, std::slice::from_ref(new_frag), &changed, + &deps, ); } } @@ -178,17 +281,22 @@ impl Transaction { pub(super) fn prune_overlay_stale_fields_from_indices( indices: &mut [IndexMetadata], groups: &[RewriteGroup], + schema: &Schema, ) { + // Resolved once for all groups, as in `prune_merge_rewritten_fields_from_indices`: + // this is the commit path, and the per-index form costs a schema walk and two + // HashSet allocations for every (group, index) pair. + let deps = Self::index_dependent_field_ids_for_each(indices, schema); for group in groups { // field id -> newest overlay committed_version supplying that field - let mut overlaid_field_versions: HashMap = HashMap::new(); + let mut overlaid_field_versions: HashMap = HashMap::new(); for old_frag in &group.old_fragments { for overlay in &old_frag.overlays { for &field_id in overlay.data_file.fields.iter() { - if field_id < 0 { - // Tombstoned (obsolete) overlay field: supplies nothing. + // Tombstoned (obsolete) overlay fields (< 0) supply nothing. + let Ok(field_id) = u32::try_from(field_id) else { continue; - } + }; let entry = overlaid_field_versions.entry(field_id).or_insert(0); *entry = (*entry).max(overlay.committed_version); } @@ -203,8 +311,12 @@ impl Transaction { .iter() .map(|f| f.id as u32) .collect::>(); - for index in indices.iter_mut() { - let is_stale = index.fields.iter().any(|field_id| { + for (index, dependent_field_ids) in indices.iter_mut().zip(deps.iter()) { + // A covered (included) field an overlay supplied makes the index just as + // stale as an indexed field would -- the index storage still holds the + // pre-overlay copy of that column. `deps` expands covered structs to their + // leaf subtree, since the overlay lists leaf ids. + let is_stale = dependent_field_ids.iter().any(|field_id| { overlaid_field_versions .get(field_id) .is_some_and(|&overlay_version| overlay_version > index.dataset_version) @@ -218,6 +330,87 @@ impl Transaction { } } + /// Reject a commit that would drop, rename, retype, or otherwise change the subtree of + /// a field an index *covers* (an "included" column). A covered field's physical payload + /// schema is fixed at index build time, so an index whose key survives but whose covered + /// column's subtree changed would keep emitting the old payload schema while covered ANN + /// queries declare the new one. This fires for two commit operations: + /// - `Project`, which can drop/rename/retype a covered field (public drop/alter APIs + /// preflight this, but a raw `Operation::Project` reaches `build_manifest` directly); + /// - `Merge` (how `add_columns` commits), which can grow a covered *struct* by adding a + /// child -- even an AllNulls, metadata-only child that writes no data file, so the + /// file-path-keyed coverage prune never sees it. + /// + /// `old_schema` is the pre-commit schema (needed to detect renames/retypes/child-adds, + /// which keep the covered field's id). `op` names the operation for the error message. + pub(super) fn reject_covered_field_subtree_change( + indices: &[IndexMetadata], + old_schema: Option<&Schema>, + new_schema: &Schema, + op: &str, + ) -> Result<()> { + let new_ids: HashSet = new_schema.fields_pre_order().map(|f| f.id).collect(); + for index in indices { + // Skip only if a KEYED field left the schema: `retain_relevant_indices` drops + // the whole index once any of `fields` is gone, so a dropped key field means + // the index is on its way out anyway and there is no dangling covering to + // protect. Checking the *keyed* prefix specifically (not all of `fields`) + // matters because `fields` also lists the covering suffix (see + // `IndexMetadata::covering_fields`): a commit that drops *only* a covering + // field, leaving the key intact, must NOT be skipped here. If it were, this + // guard would let the drop through uncontested, and `retain_relevant_indices` + // -- which keys its own retention on the very same full `fields` list -- would + // then silently delete the entire index (not just desync its covering) with no + // error. Below, that same dropped covering id is caught by + // `covered_field_subtree_changed`'s "removed" case and rejected instead. + // (`FLAG_COVERED_INDEX_METADATA` fences pre-covering builds off the whole + // dataset, so a declaration erased by an old writer -- once this guard's + // other concern -- cannot arise.) + if !index.keyed_fields().iter().all(|id| new_ids.contains(id)) { + continue; + } + for &covered_id in index.covering_fields.iter() { + let changed = match old_schema { + Some(old) => Self::covered_field_subtree_changed(old, new_schema, covered_id), + // Without the pre-commit schema we can only detect a drop. + None => new_schema.field_by_id(covered_id).is_none(), + }; + if changed { + return Err(Error::invalid_input(format!( + "{op} would drop or alter covered (included) field id {covered_id} \ + (its subtree changed), still used by index '{}'. Drop the index with \ + drop_index() before changing the column.", + index.name + ))); + } + } + } + Ok(()) + } + + /// Whether a covered ("included") field's subtree differs between `old_schema` and + /// `new_schema` -- dropped, renamed, retyped, a nullability flip, or any child change. + /// Such a change means the index's stored physical payload no longer matches the + /// schema a query declares. `data_type()` is recursive for structs, so it also covers + /// a child's name/type/nullability change. + pub fn covered_field_subtree_changed( + old_schema: &Schema, + new_schema: &Schema, + covered_id: i32, + ) -> bool { + let Some(old) = old_schema.field_by_id(covered_id) else { + return false; + }; + match new_schema.field_by_id(covered_id) { + None => true, + Some(new) => { + old.name != new.name + || old.nullable != new.nullable + || old.data_type() != new.data_type() + } + } + } + pub(crate) fn retain_relevant_indices( indices: &mut Vec, schema: &Schema, @@ -451,6 +644,8 @@ impl Transaction { mod tests { use super::*; use crate::transaction::test_support::overlay_with_field; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use lance_core::datatypes::Schema as LanceSchema; use uuid::Uuid; #[test] @@ -942,6 +1137,13 @@ mod tests { // Post-remap state: every index already covers the new fragment (7). let covering = || Some(RoaringBitmap::from_iter([7u32])); + // Indexes an un-overlaid field (3) but *covers* the overlaid field 1 as an + // included column; built (v2) before the overlay, so its stored copy of + // field 1 is stale and the rewritten fragment must be dropped. + let mut covered_stale = create_test_index("covered_stale", 3, 2, covering(), false); + // `covering_fields` is always the trailing entries of `fields`. + covered_stale.fields = vec![3, 1]; + covered_stale.covering_fields = vec![1]; let mut indices = vec![ // Stale: covers the overlaid field 1, built (v2) before the overlay. create_test_index("stale", 1, 2, covering(), false), @@ -950,9 +1152,16 @@ mod tests { create_test_index("fresh", 1, 5, covering(), false), // Unrelated: covers field 2, which the overlay never touched. create_test_index("unrelated", 2, 2, covering(), false), + covered_stale, ]; - Transaction::prune_overlay_stale_fields_from_indices(&mut indices, &groups); + // Flat field ids with no struct nesting: an empty schema makes subtree + // expansion a no-op (each id maps to itself), exercising the id-level logic. + Transaction::prune_overlay_stale_fields_from_indices( + &mut indices, + &groups, + &LanceSchema::default(), + ); assert!( !indices[0].fragment_bitmap.as_ref().unwrap().contains(7), @@ -966,6 +1175,207 @@ mod tests { indices[2].fragment_bitmap.as_ref().unwrap().contains(7), "an index on an un-overlaid field is unaffected" ); + assert!( + !indices[3].fragment_bitmap.as_ref().unwrap().contains(7), + "an index whose *covered* field was overlaid is stale too" + ); + } + + /// `covering_fields` records a covered struct's parent id, but a modified data file + /// lists leaf ids. Schema-aware pruning must expand the covered struct to its subtree + /// so an update to a subfield still invalidates the fragment's coverage. + #[test] + fn test_prune_updated_fields_expands_covered_struct_subtree() { + use arrow_schema::Fields as ArrowFields; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ArrowField::new("x", DataType::Int32, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let s_id = schema.field("s").unwrap().id; + let a_id = u32::try_from(schema.field("s.a").unwrap().id).unwrap(); + let x_id = schema.field("x").unwrap().id; + + // Index keyed on scalar `x`, covering the whole struct `s`, over fragment 0. + let mut idx = create_test_index( + "cov", + x_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + false, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![x_id, s_id]; + idx.covering_fields = vec![s_id]; + + // Updating leaf `s.a` prunes the covered struct's fragment: the subtree + // expansion bridges the parent-id-vs-leaf-id mismatch that exact-id matching + // would miss. + let mut indices = vec![idx]; + Transaction::prune_updated_fields_from_indices( + &mut indices, + &[Fragment::new(0)], + &[a_id], + &schema, + ); + assert!( + !indices[0].fragment_bitmap.as_ref().unwrap().contains(0), + "updating a covered struct's subfield must drop the fragment from coverage" + ); + } + + /// A Project that drops or renames a covered (included) column must be rejected at + /// the commit boundary, even when the index's key column survives. + #[test] + fn test_reject_projected_covered_fields() { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Utf8, true), + ArrowField::new("other", DataType::Int32, true), + ]); + let old_schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = old_schema.field("vec").unwrap().id; + let meta_id = old_schema.field("meta").unwrap().id; + + // Index keyed on `vec`, covering `meta`. + let mut idx = create_test_index( + "cov", + vec_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + true, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![vec_id, meta_id]; + idx.covering_fields = vec![meta_id]; + let indices = vec![idx]; + + // Keeping every field -> Ok. + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &old_schema, + "Project" + ) + .is_ok() + ); + + // Dropping the covered `meta` while keeping the key -> rejected. + let dropped = old_schema.project(&["vec", "other"]).unwrap(); + let err = Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &dropped, + "Project", + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered"), + "expected a covered-subtree rejection, got: {err}" + ); + + // Dropping the key too (the index would be dropped) -> not our concern. + let drop_all = old_schema.project(&["other"]).unwrap(); + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &drop_all, + "Project" + ) + .is_ok() + ); + + // Renaming the covered field (id stays, name changes) -> rejected. + let mut renamed = old_schema.clone(); + renamed.field_by_id_mut(meta_id).unwrap().name = "meta_renamed".to_string(); + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &renamed, + "Project" + ) + .is_err(), + "renaming a covered field must be rejected" + ); + + // Flipping the covered field's nullability (id + name stay) -> rejected. + let mut retyped_null = old_schema.clone(); + retyped_null.field_by_id_mut(meta_id).unwrap().nullable = false; + assert!( + Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &retyped_null, + "Project" + ) + .is_err(), + "a covered field nullability change must be rejected" + ); + } + + /// Growing a covered *struct*'s subtree (adding a child) is what `add_columns`/Merge + /// does; the covered field id is unchanged but its `data_type()` differs, so the commit + /// boundary must reject it just like the Project drop/alter cases above. + #[test] + fn test_reject_covered_struct_child_add() { + let child = ArrowField::new("a", DataType::Int32, false); + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Struct(vec![child.clone()].into()), true), + ]); + let old_schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = old_schema.field("vec").unwrap().id; + let meta_id = old_schema.field("meta").unwrap().id; + + let mut idx = create_test_index( + "cov", + vec_id, + 1, + Some(RoaringBitmap::from_iter([0u32])), + true, + ); + // `covering_fields` is always the trailing entries of `fields`. + idx.fields = vec![vec_id, meta_id]; + idx.covering_fields = vec![meta_id]; // covers the struct's parent id + let indices = vec![idx]; + + // New schema grows `meta` with a second child -> the struct's data_type changes. + let grown_arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new( + "meta", + DataType::Struct(vec![child, ArrowField::new("b", DataType::Int32, true)].into()), + true, + ), + ]); + // Preorder id assignment gives `meta` the same id (1) in both schemas, so the + // covered id resolves to the struct on each side and only its data_type differs. + let grown = LanceSchema::try_from(&grown_arrow).unwrap(); + assert_eq!(grown.field("meta").unwrap().id, meta_id); + + let err = Transaction::reject_covered_field_subtree_change( + &indices, + Some(&old_schema), + &grown, + "add_columns (Merge)", + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered") + && err.to_string().contains("add_columns"), + "expected a covered-struct-growth rejection, got: {err}" + ); } // Helper functions for retain_relevant_indices tests diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 6f216f78c06..bbc16c9cd8e 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -697,6 +697,7 @@ impl Transaction { &mut final_indices, updated_fragments, fields_modified, + &schema, ); let mut new_fragments = @@ -841,7 +842,7 @@ impl Transaction { // base data. Any index older than one of those overlays was built on // the pre-overlay values, so drop the rewritten fragment from its // coverage to keep it from serving stale values. - Self::prune_overlay_stale_fields_from_indices(&mut final_indices, groups); + Self::prune_overlay_stale_fields_from_indices(&mut final_indices, groups, &schema); if let Some(frag_reuse_index) = frag_reuse_index { final_indices.retain(|idx| idx.name != frag_reuse_index.name); @@ -870,6 +871,45 @@ impl Transaction { new_index.validate_covering_fields()?; } final_indices.extend(new_indices.clone()); + // All delta segments of one logical index (same name) must declare + // the same covering columns: the covered read path derives its + // output schema from one delta but executes all of them, so + // committing a mixed set would fail every query on the index. + // Scoped to names this operation adds, so a commit untouched by + // covering cannot be blocked by pre-existing state. + for new_index in new_indices { + // A declared covering field must resolve to a TOP-LEVEL field of + // the schema being committed: the read path resolves the id to its + // bare name, so a nested id (e.g. a struct child) would silently + // bind to a same-named top-level column, and a nonexistent id + // fails every later query at planning time ("index metadata and + // schema are inconsistent"). Reject at the commit boundary. + for id in &new_index.covering_fields { + if !schema.fields.iter().any(|field| field.id == *id) { + return Err(Error::invalid_input(format!( + "CreateIndex: index '{}' (segment {}) declares covering \ + field id {}, which is not a top-level field of the \ + schema (covering columns must be top-level)", + new_index.name, new_index.uuid, id + ))); + } + } + if let Some(conflict) = final_indices.iter().find(|idx| { + idx.name == new_index.name + && idx.covering_fields != new_index.covering_fields + }) { + return Err(Error::invalid_input(format!( + "CreateIndex: delta segment {} of index '{}' declares covering \ + fields {:?}, but delta segment {} declares {:?}; all segments of \ + one index must declare the same covering columns", + new_index.uuid, + new_index.name, + new_index.covering_fields, + conflict.uuid, + conflict.covering_fields + ))); + } + } } Operation::ReserveFragments { .. } | Operation::UpdateConfig { .. } => { final_fragments.extend(maybe_existing_fragments?.clone()); @@ -906,6 +946,18 @@ impl Transaction { } final_fragments.extend(merged_fragments); + // `add_columns` commits as a Merge; growing a *covered* field's subtree + // (e.g. an AllNulls child under a covered struct, which writes no data file + // and so is invisible to the file-path-keyed prune below) would leave the + // index emitting the old payload schema while covered queries declare the + // new one. Reject at the commit boundary, mirroring the Project guard. + Self::reject_covered_field_subtree_change( + &final_indices, + current_manifest.map(|m| &m.schema), + &schema, + "add_columns (Merge)", + )?; + // A Merge can rewrite a column's data file in place; the field stays // in the schema, so the index is retained -- prune its now-stale // entries for the rewritten fragments. @@ -913,6 +965,7 @@ impl Transaction { &mut final_indices, existing_fragments, fragments, + &schema, ); // Some fields that have indices may have been removed, so we should @@ -936,6 +989,16 @@ impl Transaction { }); } + // A Project that drops or renames a *covered* field would leave dangling + // index metadata whose storage still emits the column -- reject it before + // dropping/retaining any indices. + Self::reject_covered_field_subtree_change( + &final_indices, + current_manifest.map(|m| &m.schema), + &schema, + "Project", + )?; + // Some fields that have indices may have been removed, so we should // remove those indices as well. Self::retain_relevant_indices(&mut final_indices, &schema, &final_fragments) @@ -1162,6 +1225,7 @@ impl Transaction { &mut final_indices, &modified_fragments, &replaced_fields, + &schema, ); } Operation::DataOverlay { groups } => { @@ -1573,6 +1637,95 @@ mod tests { ) } + /// A stable-row-id `RewriteRows` update that modifies ONLY a covering + /// (included) column must not re-extend a covered index's fragment + /// coverage onto the moved fragment: the index still holds the pre-update + /// materialized value, so reusing it would serve a stale covered result. + /// The modified field is in `covering_fields`, not `fields`. + /// + /// Both spellings of "the covering column changed" must be caught. `covering_fields` + /// only ever holds TOP-LEVEL ids (index creation resolves whole columns and rejects + /// dotted names), but the caller may report either the covered struct's own id or the + /// leaf ids it actually rewrote -- `Operation::Update` is public, so + /// `fields_for_preserving_frag_bitmap` is caller-supplied and unvalidated. Comparing raw + /// ids catches only the first spelling and silently serves stale subfield values for the + /// second. + #[test] + fn test_pure_rewrite_preserves_covering_index_coverage() { + use arrow_schema::Fields as ArrowFields; + + // meta{a} => ids 0 (struct) and 1 (leaf); vector => id 2. + let arrow = ArrowSchema::new(vec![ + ArrowField::new( + "meta", + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])), + true, + ), + ArrowField::new("vector", DataType::Int32, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let meta_id = schema.field("meta").unwrap().id; + let leaf_id = schema.field("meta.a").unwrap().id; + let vector_id = schema.field("vector").unwrap().id; + + // Returns (covered index re-extended?, plain index re-extended?). + let run = |modified: &[u32]| { + let mut covering = sample_index_metadata("covering"); + // `covering_fields` is always the trailing entries of `fields`. + covering.fields = vec![vector_id, meta_id]; // top-level, as creation guarantees + covering.covering_fields = vec![meta_id]; + covering.fragment_bitmap = Some([0].into_iter().collect()); + + let mut plain = sample_index_metadata("plain"); + plain.fields = vec![vector_id]; + plain.covering_fields = vec![]; + plain.fragment_bitmap = Some([0].into_iter().collect()); + + let mut indices = vec![covering, plain]; + let overlaid: HashMap = HashMap::new(); + Transaction::register_pure_rewrite_rows_update_frags_in_indices( + &mut indices, + &[1], // new fragment holding the moved rows + &[0], // original fragment (covered by both indices) + modified, + &overlaid, + &schema, + ) + .unwrap(); + ( + indices[0].fragment_bitmap.as_ref().unwrap().contains(1), + indices[1].fragment_bitmap.as_ref().unwrap().contains(1), + ) + }; + + for (label, modified) in [ + ( + "covered struct reported by its own id", + vec![meta_id as u32], + ), + ( + "covered struct reported by its leaf id", + vec![leaf_id as u32], + ), + ] { + let (covered_extended, plain_extended) = run(&modified); + assert!( + !covered_extended, + "covered index must not re-extend coverage onto a fragment whose covering \ + column was rewritten ({label})" + ); + assert!( + plain_extended, + "index that neither indexes nor covers the modified field still reuses its \ + entries ({label})" + ); + } + } + #[test] fn test_create_index_build_manifest_keeps_unremoved_same_name_indices() { let manifest = sample_manifest(); @@ -1649,6 +1802,251 @@ mod tests { ); } + /// A raw CreateIndex commit whose new index declares a covering field id missing + /// from the schema must be rejected at the commit boundary -- otherwise every later + /// query on the index fails at planning time with an inconsistent-metadata error. + /// + /// `covering_fields` must also be a suffix of `fields` (`fields = [0, 9999]` here), + /// or `validate_covering_fields`'s cardinality check rejects the index earlier, for + /// an unrelated reason, before this test's intended "unresolvable in schema" check + /// ever runs. + #[test] + fn test_create_index_build_manifest_rejects_unresolvable_covered_field() { + let manifest = sample_manifest(); + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, 9999]; + covered.covering_fields = vec![9999]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("9999") && err.to_string().contains("top-level"), + "an unresolvable covered field id must be rejected at commit: {err}" + ); + } + + /// A covering declaration must reference a TOP-LEVEL field: `covering_fields` + /// resolution on the read path goes id -> name -> arrow field by bare name, so a + /// nested field id (e.g. a struct child) would silently bind to a same-named + /// top-level column instead of erroring. The create path rejects dotted names; + /// raw commits must reject nested ids for the same reason. + #[test] + fn test_create_index_build_manifest_rejects_nested_covered_field() { + use arrow_schema::Fields as ArrowFields; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "id", + DataType::Int64, + true, + )])), + true, + ), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let nested_id = schema.field("s.id").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, nested_id]; + covered.covering_fields = vec![nested_id]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("top-level"), + "a nested covered field id must be rejected at commit: {err}" + ); + } + + /// A raw `Operation::Project` that drops ONLY a covering (included) field, keeping the + /// index's key field intact, must be rejected at the commit boundary. `fields` lists the + /// covering suffix alongside the key (see `IndexMetadata::covering_fields`), so + /// `retain_relevant_indices` -- which requires every entry of `fields` to survive -- + /// would otherwise auto-prune the WHOLE index the moment the commit went through, with + /// no error: a working vector index silently disappearing rather than merely desyncing + /// its covering payload. + #[test] + fn test_project_build_manifest_rejects_dropped_covering_field() { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("meta", DataType::Utf8, true), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let vec_id = schema.field("vec").unwrap().id; + let meta_id = schema.field("meta").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![vec_id, meta_id]; + covered.covering_fields = vec![meta_id]; + + let projected_schema = manifest.schema.project(&["vec"]).unwrap(); + let transaction = Transaction::new( + manifest.version, + Operation::Project { + schema: projected_schema, + preserves_nullability: true, + }, + None, + ); + let err = transaction + .build_manifest( + Some(&manifest), + vec![covered], + "txn", + &default_build_config(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("drop or alter covered"), + "dropping a covering field while its index's key survives must be rejected \ + instead of silently letting the whole index be auto-pruned: {err}" + ); + } + + /// All delta segments of one logical index must declare the same covering columns: + /// the covered read path derives its output schema from one delta but executes all + /// of them, so a mixed set fails every query on the index at execution time. + #[test] + fn test_create_index_build_manifest_rejects_mixed_covering_deltas() { + // A second top-level field is needed so the covered index can declare it as a + // trailing, schema-resolvable `covering_fields` entry (`sample_manifest`'s + // single-field schema has no field left over once field 0 is keyed on). + let arrow = ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("vector", DataType::Int32, false), + ]); + let schema = LanceSchema::try_from(&arrow).unwrap(); + let covered_field_id = schema.field("vector").unwrap().id; + let manifest = Manifest::new( + schema, + Arc::new(vec![Fragment::new(0)]), + DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + let mut covered = sample_index_metadata("vector_idx"); + covered.fields = vec![0, covered_field_id]; + covered.covering_fields = vec![covered_field_id]; + let plain = sample_index_metadata("vector_idx"); + + // Adding a non-covered delta to an existing covered index is rejected. + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![plain.clone()], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest( + Some(&manifest), + vec![covered.clone()], + "txn", + &default_build_config(), + ) + .unwrap_err(); + assert!( + err.to_string().contains("covering"), + "mixed covering across deltas must be rejected: {err}" + ); + + // A mixed covering set within one commit's own new segments is rejected too. + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered.clone(), plain.clone()], + removed_indices: vec![], + }, + None, + ); + let err = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap_err(); + assert!( + err.to_string().contains("covering"), + "mixed covering within one commit must be rejected: {err}" + ); + + // Removing the disagreeing delta in the same commit leaves a homogeneous + // final state, which is fine. + let mut covered2 = sample_index_metadata("vector_idx"); + covered2.fields = vec![0, covered_field_id]; + covered2.covering_fields = vec![covered_field_id]; + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![covered2], + removed_indices: vec![plain.clone()], + }, + None, + ); + let (_, final_indices) = transaction + .build_manifest( + Some(&manifest), + vec![covered.clone(), plain], + "txn", + &default_build_config(), + ) + .unwrap(); + assert_eq!(final_indices.len(), 2); + assert!( + final_indices + .iter() + .all(|idx| idx.covering_fields == vec![covered_field_id]), + "the homogeneous replacement set should commit" + ); + + // An index with a different name is unaffected. + let other = sample_index_metadata("other_idx"); + let transaction = Transaction::new( + manifest.version, + Operation::CreateIndex { + new_indices: vec![other], + removed_indices: vec![], + }, + None, + ); + transaction + .build_manifest( + Some(&manifest), + vec![covered], + "txn", + &default_build_config(), + ) + .unwrap(); + } + #[test] fn test_update_build_manifest_replaces_and_removes_fragments() { let manifest = sample_manifest_with_fragments(0..5); diff --git a/rust/lance/src/dataset/optimize/remapping.rs b/rust/lance/src/dataset/optimize/remapping.rs index 8a9c8898cb6..065612ed862 100644 --- a/rust/lance/src/dataset/optimize/remapping.rs +++ b/rust/lance/src/dataset/optimize/remapping.rs @@ -307,16 +307,23 @@ async fn remap_index(dataset: &mut Dataset, index_id: &Uuid) -> Result<()> { // does not incorporate overlays committed after the source index was built. // Exclude those fragments so queries scan their current values instead. if let Some(fragment_bitmap) = &mut bitmap_after_remap { + // Fields whose value this index materializes: its indexed key fields plus every + // covered (included) column expanded to its leaf subtree (an overlay lists leaf + // ids, while `covering_fields` records a parent struct id). + let relevant_field_ids = + Transaction::index_dependent_leaf_ids(&curr_index_meta, dataset.schema()); for fragment in dataset.manifest.fragments.iter() { - let has_newer_indexed_overlay = fragment.overlays.iter().any(|overlay| { + // An overlay refreshing an indexed or covered field makes the index storage + // stale for that fragment: it still holds the pre-overlay copy of the value. + let has_newer_relevant_overlay = fragment.overlays.iter().any(|overlay| { overlay.committed_version > curr_index_meta.dataset_version && overlay .data_file .fields .iter() - .any(|field_id| curr_index_meta.fields.contains(field_id)) + .any(|field_id| relevant_field_ids.contains(field_id)) }); - if has_newer_indexed_overlay { + if has_newer_relevant_overlay { fragment_bitmap.remove(fragment.id as u32); } } diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index af0e1ab2622..b93f90f3bd6 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -793,37 +793,102 @@ pub(super) async fn alter_columns( new_schema.validate()?; new_schema.verify_primary_key()?; - // If any column being cast has an attached index, fail fast. Cast operations - // rewrite the underlying column data and silently invalidate any index on the - // affected column(s). The current behavior is to drop such indices without - // warning, which has caused production incidents where vector search silently - // regressed to brute-force scan. We require users to explicitly drop the - // index before altering the column type, so the action is never silent. That - // includes an index this build has no reader for: the cast reassigns the - // field id, so carrying it forward is impossible and staying quiet about it - // is the silent drop this guard exists to abolish. - if !cast_fields.is_empty() { + // Fail fast on alterations that would invalidate an index or desync its + // covering data: + // * CAST of an indexed *key* column -- cast rewrites the column data and + // silently invalidates the index. This has caused production incidents + // where vector search regressed to a brute-force scan. (Indices are + // keyed by field id, so a plain rename or nullability change of a key + // column is safe and stays allowed.) + // * ANY alter (rename, cast, or nullability) of a covering ("included") + // column. `fields` lists the covering suffix right alongside the key + // (see `IndexMetadata::covering_fields`), so the auto-prune + // (`retain_relevant_indices`) DOES see a covering column's id. + // + // For a CAST this guard is load-bearing: a cast reassigns the column a + // fresh field id, so the old covering id disappears from the schema and + // auto-prune -- which requires every id in `fields` to survive -- would + // delete the WHOLE index, not just leave its covering stale. + // + // For a plain rename or nullability flip the id survives, so auto-prune + // never fires and the index is left declaring covering that its storage + // no longer matches. That is NOT silent corruption: the read path + // intersects the declaration with each segment's proven physical + // capability (`common_physical_covering`, which compares name, type, + // nullability and metadata), so a mismatch withdraws covering and the + // query falls back to a base-table take -- correct results, minus the + // elision. The guard is conservative here rather than load-bearing, and + // is kept so a lost elision surfaces as an explicit error instead of a + // silent slowdown. Relaxing it to degrade instead of reject is a tracked + // follow-up. + // + // Both cases include an index this build has no reader for -- hence + // `load_all_indices`. A cast reassigns the field id, so carrying such an index + // forward is impossible, and staying quiet about it is the silent drop these + // guards exist to abolish. + // Reject rather than let the index quietly change behaviour; the cast case + // above is the one that would otherwise destroy the index outright. + let cast_ids: Vec = alterations + .iter() + .filter(|a| a.data_type.is_some()) + .filter_map(|a| dataset.schema().field(&a.path).map(|f| f.id)) + .collect(); + let altered: Vec<(&str, i32)> = alterations + .iter() + .filter(|a| a.rename.is_some() || a.data_type.is_some() || a.nullable.is_some()) + .filter_map(|a| { + dataset + .schema() + .field(&a.path) + .map(|f| (a.path.as_str(), f.id)) + }) + .collect(); + if !altered.is_empty() { let indices = load_all_indices(dataset).await?; - let affected: Vec<&lance_table::format::IndexMetadata> = indices - .iter() - .filter(|idx| { - cast_fields - .iter() - .any(|(old, _)| idx.fields.contains(&old.id)) + // A covered column is stored as a whole subtree, so altering a *descendant* of a + // covered field (e.g. a subfield of a covered struct) invalidates the covering + // too; expand each altered field to its ancestors so the covered ancestor's id is + // caught. The expansion depends only on the schema, so do it once per altered + // field here rather than per (index x altered field) below. + let altered: Vec<(&str, i32, Vec)> = altered + .into_iter() + .map(|(name, id)| { + let ancestors = + field_ids_with_ancestors(dataset.schema(), std::slice::from_ref(&id)); + (name, id, ancestors) }) .collect(); + // An index is broken if a covering column of it is altered in any way, or if one + // of its indexed key columns is being cast. The cast check stays on the field's + // own id (only a cast of the indexed key column rewrites index data). + let is_broken = |idx: &lance_table::format::IndexMetadata| { + altered.iter().any(|(_, _, ancestors)| { + ancestors + .iter() + .any(|ancestor| idx.covering_fields.contains(ancestor)) + }) || cast_ids.iter().any(|id| idx.fields.contains(id)) + }; + let affected: Vec<&lance_table::format::IndexMetadata> = + indices.iter().filter(|idx| is_broken(idx)).collect(); if !affected.is_empty() { - let affected_cols: Vec = cast_fields + let affected_cols: Vec = altered .iter() - .filter(|(old, _)| affected.iter().any(|i| i.fields.contains(&old.id))) - .map(|(old, _)| old.name.clone()) + .filter(|(_, id, ancestors)| { + affected.iter().any(|i| { + ancestors + .iter() + .any(|ancestor| i.covering_fields.contains(ancestor)) + || (cast_ids.contains(id) && i.fields.contains(id)) + }) + }) + .map(|(name, _, _)| name.to_string()) .collect(); let affected_idx_names: Vec = affected.iter().map(|i| i.name.clone()).collect(); return Err(Error::invalid_input(format!( - "Cannot cast column(s) [{}] to a new type: they have {} index(es) \ - attached: [{}]. Cast rewrites column data and invalidates any index \ - on the affected column(s). Drop the index(es) with drop_index() \ - before altering, then recreate them after the cast completes.", + "Cannot alter column(s) [{}]: they are used by {} index(es) [{}] -- as a \ + covering (included) column (any change), or as an indexed key column being \ + cast (which rewrites the data and invalidates the index). Drop the index(es) \ + with drop_index() before altering, then recreate them.", affected_cols.join(", "), affected.len(), affected_idx_names.join(", "), @@ -1017,6 +1082,23 @@ pub(super) async fn alter_columns( Ok(()) } +/// Expand `field_ids` to also include every ancestor field id (root -> field). +/// A covering column is stored as a whole subtree, so touching a *descendant* of a +/// covered field (e.g. a subfield of a covered struct) still invalidates the covering +/// even though `covering_fields` records only the ancestor's id. Adding each targeted +/// field's ancestors makes the ancestor's id appear in the set so the exact-membership +/// checks catch it. +fn field_ids_with_ancestors(schema: &Schema, field_ids: &[i32]) -> Vec { + let mut expanded = Vec::with_capacity(field_ids.len()); + for &id in field_ids { + match schema.field_ancestry_by_id(id) { + Some(ancestry) => expanded.extend(ancestry.iter().map(|field| field.id)), + None => expanded.push(id), + } + } + expanded +} + /// Remove columns from the dataset. /// /// This is a metadata-only operation and does not remove the data from the @@ -1045,6 +1127,63 @@ pub(super) async fn drop_columns(dataset: &mut Dataset, columns: &[&str]) -> Res )); } + // Fail fast if any column being dropped is a covering ("included") column of + // an index, while the index's key survives. Dropping the key too is fine -- + // `retain_relevant_indices` auto-prunes the whole index once any id in `fields` + // leaves the schema, and `fields` lists the covering suffix right alongside the + // key (see `IndexMetadata::covering_fields`), so a key+covering drop cleanly + // removes the whole thing. But that is exactly why dropping *only* the covering + // column, without this guard, would not merely leave the index's storage + // referencing a missing column -- the same auto-prune check would see the + // covering id vanish from `fields` and silently delete the ENTIRE index (the key, + // its ANN structure, everything), not just desync its covering. Require the user + // to drop the index first so the action is never silent. + let drop_field_ids: Vec = columns + .iter() + .filter_map(|c| dataset.schema().field(c).map(|f| f.id)) + .collect(); + let indices = load_all_indices(dataset).await?; + // Treat a descendant of a covered (e.g. struct) column as covered too: the covering + // declaration names the struct while the drop names a leaf inside it. + let drop_field_ids = field_ids_with_ancestors(dataset.schema(), &drop_field_ids); + // An index whose *key* is also being dropped is auto-pruned by + // `retain_relevant_indices`, so its covering metadata cannot dangle -- the whole index + // goes away. Rejecting here would block the perfectly safe "drop the vector column and + // its covered payload together" case. The commit-boundary guard + // (`reject_covered_field_subtree_change`) skips these indices for the same reason, and + // this preflight must agree with it or it forbids what the commit would accept. + // + // Agreement requires testing what that guard tests: whether a KEYED field survives the + // post-drop schema. Neither half of that can be read off `drop_field_ids`. It is + // expanded *upwards* to ancestors, so dropping `s.a` lists the struct `s` that in fact + // survives; and the keyed prefix is not all of `fields`, which also carries the + // covering suffix. Testing `fields` against the expanded set got both wrong in + // opposite directions -- it rejected dropping a key alongside a strict subset of its + // covering columns (the commit accepts that: the key is gone, the index goes with it), + // and it skipped an index keyed on a struct when a leaf under that struct was dropped, + // deferring a case this preflight is meant to catch to a far less legible error. + let surviving: HashSet = new_schema.fields_pre_order().map(|f| f.id).collect(); + let affected: Vec<&lance_table::format::IndexMetadata> = indices + .iter() + .filter(|idx| idx.keyed_fields().iter().all(|id| surviving.contains(id))) + .filter(|idx| { + drop_field_ids + .iter() + .any(|id| idx.covering_fields.contains(id)) + }) + .collect(); + if !affected.is_empty() { + let affected_idx_names: Vec = affected.iter().map(|i| i.name.clone()).collect(); + return Err(Error::invalid_input(format!( + "Cannot drop column(s) [{}]: they are a covering (included) column of {} \ + index(es) [{}]. Drop the index(es) with drop_index() before dropping the \ + column(s), then recreate them.", + columns.join(", "), + affected.len(), + affected_idx_names.join(", "), + ))); + } + let transaction = Transaction::new( dataset.manifest.version, Operation::Project { @@ -3589,6 +3728,365 @@ mod test { Ok(()) } + /// The covered-column schema guards must treat a *descendant* of a covered field as + /// covered too: a covered column is stored as a whole subtree, so `covering_fields` + /// records only the ancestor's id. `field_ids_with_ancestors` expands a targeted + /// subfield to include that ancestor so the exact-membership guard catches it. + #[test] + fn test_field_ids_with_ancestors_expands_struct_subtree() { + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ArrowField::new("x", DataType::Int32, true), + ]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let s_id = schema.field("s").unwrap().id; + let a_id = schema.field("s.a").unwrap().id; + + let expanded = field_ids_with_ancestors(&schema, &[a_id]); + assert!( + expanded.contains(&s_id), + "expanding a subfield must include its covered struct ancestor's id" + ); + assert!( + expanded.contains(&a_id), + "the field's own id must be included" + ); + } + + /// Dropping a covered column TOGETHER WITH its index's key column is safe and must be + /// allowed: `retain_relevant_indices` drops the whole index once its key leaves the + /// schema, so no covering metadata can dangle. The commit-boundary guard + /// (`reject_covered_field_subtree_change`) already skips such indices; this preflight + /// must agree with it, or it forbids an operation the commit would accept. + /// + /// The `key_and_covering_subset` case is what makes the two agree in general. The + /// preflight used to skip an index only when *every* id in `fields` was being dropped, + /// while the commit guard skips as soon as a *keyed* id is. With two covering columns + /// and only one of them dropped alongside the key, those disagree: the commit accepts + /// (the key is gone, so the whole index goes) and the preflight rejected. + #[rstest] + #[case::key_and_all_covering(&["vec", "id", "extra"], vec!["keep"])] + #[case::key_and_covering_subset(&["vec", "id"], vec!["extra", "keep"])] + #[tokio::test] + async fn test_drop_covered_column_with_its_index_key_is_allowed( + #[case] to_drop: &[&str], + #[case] expected_remaining: Vec<&str>, + ) -> Result<()> { + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + use crate::index::vector::VectorIndexParams; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("extra", DataType::Int32, false), + ArrowField::new("keep", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + Arc::new(Int32Array::from_iter_values(0..nrows)), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string(), "extra".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + + // Drop the index key together with its covering column(s) in one call. + dataset.drop_columns(to_drop).await?; + + assert!( + dataset.load_indices().await?.is_empty(), + "the index must be pruned once its key column is gone" + ); + let remaining: Vec<&str> = dataset + .schema() + .fields + .iter() + .map(|f| f.name.as_str()) + .collect(); + assert_eq!(remaining, expected_remaining); + Ok(()) + } + + /// Dropping a column that an index *covers* (an "included"/covering column, + /// not the key column) must fail fast. The covering id lives in `fields`, so + /// letting the drop through would make `retain_relevant_indices` delete the + /// ENTIRE index -- key, ANN structure and all -- not merely desync its + /// covering. The user must drop the index first. + #[tokio::test] + async fn test_drop_columns_fails_on_covered_column() -> Result<()> { + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + use crate::index::vector::VectorIndexParams; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + // IVF_PQ index on `vec` that COVERS `id`. + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + let index_name = dataset.load_indices().await?[0].name.clone(); + + // Dropping the covered column must fail, naming the column, the index, + // and the remediation. + let err = dataset + .drop_columns(&["id"]) + .await + .expect_err("dropping a covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + let msg = err.to_string(); + // `[id]` (the rendered column list), not a bare `id`: the auto-generated index name + // is `vec_idx`, which contains `id` as a substring, so a bare `contains("id")` would + // pass even if the message never named the column. + assert!( + msg.contains("[id]") && msg.contains(&index_name) && msg.contains("drop_index"), + "error should mention the column, index, and remediation, got: {msg}" + ); + + // Unchanged: column still present, index still there. + assert!(dataset.schema().field("id").is_some()); + assert_eq!(dataset.load_indices().await?.len(), 1); + + // After dropping the index, the same drop succeeds. + dataset.drop_index(&index_name).await?; + dataset.drop_columns(&["id"]).await?; + assert!(dataset.schema().field("id").is_none()); + + Ok(()) + } + + /// Dropping a column that is only an INDEXED key column (not a covering + /// column) must remain allowed: lance auto-drops the dependent index. Only + /// a covering ("included") column dropped *while its index's key survives* + /// is protected by the drop guard: dropping the key removes the whole + /// index either way (via `retain_relevant_indices`), so there is nothing to + /// protect, whereas dropping just the covering column would otherwise + /// silently take the whole index down with it (see `drop_columns`). + /// Regression guard for the Python `test_drop_columns` behavior. + #[tokio::test] + async fn test_drop_indexed_noncovered_column_drops_index() -> Result<()> { + use crate::index::vector::VectorIndexParams; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + // Plain IVF_PQ index on `vec` -- NO covering columns. + let params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + assert_eq!(dataset.load_indices().await?.len(), 1); + + // Dropping the indexed key column is allowed and auto-drops the index. + dataset.drop_columns(&["vec"]).await?; + assert!(dataset.schema().field("vec").is_none()); + assert_eq!( + dataset.load_indices().await?.len(), + 0, + "dropping the indexed column should have dropped its index" + ); + + Ok(()) + } + + /// Fix 1: renaming or changing the nullability of a covered ("included") + /// column must fail fast (alongside cast) -- it would leave the index storage + /// referencing a stale name/schema and break covered queries. + #[tokio::test] + async fn test_alter_columns_rename_nullable_fail_on_covered_column() -> Result<()> { + use crate::index::vector::VectorIndexParams; + use lance_arrow::FixedSizeListArrayExt; + use lance_index::IndexType; + use lance_linalg::distance::MetricType; + use lance_testing::datagen::generate_random_array; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 64, + ), + false, + ), + ArrowField::new("id", DataType::Int32, false), + ])); + let nrows: i32 = 256; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new( + ::try_new_from_values( + generate_random_array(64 * nrows as usize), + 64, + ) + .unwrap(), + ), + Arc::new(Int32Array::from_iter_values(0..nrows)), + ], + )?; + let test_dir = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + &test_dir, + None, + ) + .await?; + + let mut params = VectorIndexParams::ivf_pq(4, 8, 8, MetricType::L2, 50); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await?; + + // Rename of the covered column must fail with the remediation hint. + let err = dataset + .alter_columns(&[ColumnAlteration::new("id".into()).rename("uid".into())]) + .await + .expect_err("rename of covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains("drop_index"), + "rename error should suggest drop_index, got: {err}" + ); + assert!( + dataset.schema().field("id").is_some(), + "rename should not have applied" + ); + + // Nullability change of the covered column must also fail. + let err = dataset + .alter_columns(&[ColumnAlteration::new("id".into()).set_nullable(true)]) + .await + .expect_err("nullability change of covered column should fail"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected Error::InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains("drop_index"), + "nullable error should suggest drop_index, got: {err}" + ); + + Ok(()) + } + #[test] fn test_is_upcast_downcast_dictionary() { use DataType::*; diff --git a/rust/lance/src/dataset/take.rs b/rust/lance/src/dataset/take.rs index ddb76a0f720..971d1d189f4 100644 --- a/rust/lance/src/dataset/take.rs +++ b/rust/lance/src/dataset/take.rs @@ -530,6 +530,19 @@ impl TakeBuilder { self } + /// Add the row id to the output *without* requesting it as a projected column. + /// + /// Naming `_rowid` in a [`ProjectionRequest`] also sets `must_add_row_offset` + /// (see `ProjectionPlan::from_schema`), which arms the "must not target deleted + /// rows" check in [`do_take_rows`] and turns a take that legitimately returns + /// fewer rows than requested into an error. It also computes a row-offset column + /// per batch. A caller that wants the ids only to realign its own output against + /// the rows that actually came back needs neither, so it uses this instead. + pub(crate) fn include_row_id(mut self) -> Self { + Arc::make_mut(&mut self.projection).include_row_id(); + self + } + pub(super) fn with_missing_row_policy(mut self, policy: MissingRowPolicy) -> Self { self.missing_row_policy = policy; self diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 7234e8baa0b..f02f1f24a3f 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -50,6 +50,7 @@ use super::{ use crate::dataset::rowids::get_row_id_index; use crate::dataset::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::dataset::utils::CapturedRowIds; +use crate::dataset::{ProjectionRequest, TakeBuilder}; use crate::index::DatasetIndexExt; use crate::{ Dataset, @@ -1520,6 +1521,337 @@ impl MergeInsertJob { self.create_full_table_joined_stream(source).await } + /// Names of `source_schema` columns this merge would update that are covered + /// ("included") by some index. Patching such a column in place makes the index's + /// materialized copy of it stale, which forces the commit to drop the whole updated + /// fragment from every index covering it; the partial-schema path therefore rewrites + /// those updates as row-moves, where it can, so that invalidation is confined to the + /// rows actually updated. See the branch in `execute_uncommitted_impl` for what that + /// does and does not buy, and `covered_move_blocker` for the cases that fall back to + /// the in-place patch. Join keys are matched, not updated, so they are excluded. + async fn covered_columns_updated(&self, source_schema: &Schema) -> Result> { + // Only the update-matched modes actually patch matched rows into fragments; + // DoNothing / Fail leave matched rows untouched (Delete is handled separately + // and never reaches this path). Without this gate, a partial source carrying a + // covered payload under the default DoNothing + InsertAll would be wrongly + // flagged and rejected even though no covered column is ever patched in place. + if !matches!( + self.params.when_matched, + WhenMatched::UpdateAll | WhenMatched::UpdateIf(_) | WhenMatched::UpdateIfExpr(_) + ) { + return Ok(Vec::new()); + } + let indices = self.dataset.load_indices().await?; + let covered: HashSet = indices + .iter() + .flat_map(|idx| idx.covering_fields.iter().copied()) + .collect(); + if covered.is_empty() { + return Ok(Vec::new()); + } + let target = self.dataset.schema(); + let keys: HashSet<&str> = self.params.on.iter().map(|s| s.as_str()).collect(); + let mut updated = Vec::new(); + for field in source_schema.fields() { + let name = field.name(); + if keys.contains(name.as_str()) { + continue; + } + if let Some(tf) = target.field(name) + && covered.contains(&tf.id) + { + updated.push(name.clone()); + } + } + Ok(updated) + } + + /// Why `rewrite_covered_update_as_move` cannot handle this merge, if it cannot. + /// + /// Every case below is a limitation of the *move*, never of the operation: the in-place + /// patch handles all of them, and it is correct for covered columns (see the caller -- + /// the suffix-subset shape already drops the fragment from the index, so no stale copy + /// is reachable). The move is an optimisation, and an optimisation must not be able to + /// fail an operation that would otherwise succeed, so the caller logs this reason and + /// patches in place, forfeiting only confinement. Returns the reason so the log can name + /// it; the phrasing is a "because ..." clause. + fn covered_move_blocker(&self, source_schema: &Schema) -> Option { + // The move would corrupt row identity. Moving a row to a new fragment must carry its + // stable row id along, and the full-schema `RewriteRows` path does exactly that -- it + // rechunks the merger's captured `row_id_sequence` onto the new fragments and stamps + // `row_id_meta` (see the `updating_row_ids` block in `execute_uncommitted_impl`). + // `rewrite_covered_update_as_move` has no such step, so the commit would reach + // `Transaction::assign_row_ids` and mint *fresh* ids for every moved row, silently + // breaking the one guarantee stable row ids exist to provide. The in-place patch moves + // no rows at all, so it preserves those ids by construction -- strictly safer here than + // the move, not less safe. + if self.dataset.manifest.uses_stable_row_ids() { + return Some( + "the move cannot carry stable row ids onto the new fragments (it has no \ + `row_id_sequence` rechunk step), while the in-place patch preserves them by \ + moving no rows" + .into(), + ); + } + // The move would delete rows the user never touched. With `insert_not_matched`, + // `Merger::execute_batch` pushes a *second* batch holding the unmatched (`left_only`) + // rows, projected with the row-address column appended. Those rows come from the left + // side of an outer join, so their `_rowaddr` is NULL. The move extracts addresses with + // `batch[ROW_ADDR].as_primitive::().values()`, which reads the raw value + // buffer and *ignores the null mask* -- so every inserted row would contribute a bogus + // address (in practice 0), used both to take the wrong row and, via `moved_addrs`, to + // tombstone it. Never entering the move prevents that as completely as refusing the + // merge did, and `update_fragments` already routes the NULL-`_rowaddr` group to + // `handle_new_fragments`. + if self.params.insert_not_matched { + return Some( + "the move reads row addresses through the null mask and would tombstone \ + unrelated rows for the unmatched (inserted) group, which arrives with a NULL \ + `_rowaddr`" + .into(), + ); + } + // The move re-reads full rows via a take, which materializes legacy (v1) blob columns + // as description structs, not binary -- writing those back would corrupt the blob. The + // full-schema path reads them as binary via the scan provider (`SomeBlobsBinary`); the + // take path has no equivalent. The in-place patch never reads the blob column at all. + // + // Scanning the *whole dataset schema* is deliberate, not sloppy. Do not narrow this to + // the covered or updated columns: the take reads every column, so an unrelated v1 blob + // elsewhere in the schema is corrupted just the same. In fact a covered column can + // never itself be a blob -- index creation rejects blob covering columns outright + // (`crate::index::vector`, the `covering_columns` validation) -- so this guard is + // entirely about the *collateral* columns the full-row take drags in. Narrowing it + // would therefore disable it completely and reintroduce the corruption. + if self + .dataset + .schema() + .fields_pre_order() + .any(|f| f.is_blob() && !f.is_blob_v2()) + { + return Some( + "the move's full-row take reads legacy (v1) blob columns as description \ + structs rather than binary, which would corrupt them on write back" + .into(), + ); + } + // The move overlays each source column onto a full-row take by REPLACING the whole + // column, which needs the source column's Arrow type to equal the target's exactly: + // `replace_column_by_name` keeps the target schema and revalidates through + // `RecordBatch::try_new`. The partial-schema contract does not promise that. A source + // is admitted as a *subschema*, and `compare_fields` recurses with the same options, + // so a column may legally differ in its struct children, in nested nullability, or in + // nested field metadata -- all three are part of `DataType` and none is compared by + // the check that admitted the source. Note a field's OWN nullability is not part of + // its `DataType`, so the ordinary "nullable source, non-null target" partial update + // does not land here. + let target_arrow = Schema::from(self.dataset.schema()); + let type_mismatched = source_schema + .fields() + .iter() + .filter(|f| { + target_arrow + .field_with_name(f.name()) + .is_ok_and(|tf| tf.data_type() != f.data_type()) + }) + .map(|f| f.name().clone()) + .collect::>(); + if !type_mismatched.is_empty() { + return Some(format!( + "the move replaces whole columns and source column(s) {type_mismatched:?} do \ + not have exactly the target's type (struct children, nested nullability or \ + nested metadata may differ)" + )); + } + None + } + + /// Rewrite a partial-schema covered-column update as a row-move: read the full + /// existing rows for the touched addresses, overlay the updated columns, write + /// them as new fragments, and tombstone the originals. Only the moved rows leave + /// the index's reach; the rest of the fragment keeps its coverage, which is the + /// point of the move (a recall/cost property, not a correctness one -- see the + /// caller for why). + /// + /// `source` is the merger's partial output stream (`_rowaddr` + updated cols); it is consumed + /// in bounded batches so peak memory is one batch of full rows rather than the whole update. + /// Mirrors `dataset.update()`'s streamed `RewriteRows` move (`update.rs::execute_impl`). + /// Returns the resulting `RewriteRows` operation together with the moved rows' addresses, so + /// the caller can publish them as `affected_rows` and let a concurrent delete/update rebase + /// instead of hard-conflicting. + /// + /// Non-stable-row-id datasets only. The caller **falls back to the in-place path** for + /// stable-row-id datasets, insert-carrying combinations, legacy v1 blob columns, and + /// partial struct-subschema sources — it does not reject them. That fallback branch is + /// live and load-bearing, not dead code: the in-place path is correct for covered columns + /// (it drops the updated fragment from the index via `prune_updated_fields_from_indices`), + /// so this move is an optimisation that confines the bitmap damage, and an optimisation + /// must not be able to fail the operation. Restoring a hard rejection here would + /// re-break merge shapes that work today. + async fn rewrite_covered_update_as_move( + &self, + source: SendableRecordBatchStream, + target_bases_info: Option>, + ) -> Result<(Operation, RoaringTreemap)> { + let full_schema = self.dataset.schema().clone(); + // EVERY field id, leaves included (`fields_pre_order`), never just the + // top-level ids: the consumer + // (`register_pure_rewrite_rows_update_frags_in_indices`) treats this list + // as "fields whose values were updated" and compares it against each + // index's leaf-expanded dependency set -- its contract for a caller that + // cannot prove any field survived unchanged is to pass *every* field id. + // An index keyed on a nested field stores a leaf id that a top-level-only + // set does not contain, so it would escape the exclusion and have its + // bitmap wrongly extended onto the moved fragment. Both sibling row-move + // producers (`update.rs`, the insert exec) already expand to leaves for + // this exact reason. + let preserving: Vec = full_schema + .fields_pre_order() + .map(|f| f.id as u32) + .collect(); + let empty_op = |fields_bitmap: Vec| Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments: Vec::new(), + new_fragments: Vec::new(), + fields_modified: Vec::new(), + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: fields_bitmap, + update_mode: Some(RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + + // Per merger batch: take the full existing rows for its addresses and overlay the updated + // columns, streaming the result so peak memory is one batch of full rows rather than the + // whole update. On the partial-schema path the merger appends `_rowaddr` to its output + // (`Merger::output_schema`), so every batch names the originals to tombstone; the take's + // result is built from the dataset schema and never carries that column onward. + let dataset = self.dataset.clone(); + let plan = Arc::new( + ProjectionRequest::from_schema(full_schema.clone()) + .into_projection_plan(dataset.clone())?, + ); + let out_arrow = Arc::new(Schema::from(&full_schema)); + + // Record which originals to tombstone as the overlay runs. `make_rowid_capture_stream` is + // deliberately not reused here: its `AddressStyle` capture bulk-loads a `RoaringTreemap` + // with `append`, which rejects any value <= the running maximum, so it is only safe on a + // stream whose addresses are strictly ascending. That holds today, but only as an accident + // of routing: this function is reachable only via the indexed-scan join (every other + // partial source goes to v2 -- see the caller), and that scan reads in fragment order, so + // addresses arrive ascending however the user ordered the source. Swapping this `extend` + // for `append` accordingly passes the covered suite, reversed source included. `extend` is + // kept anyway: it costs nothing, and it means a future routing or planner change cannot + // turn an ordering assumption into an abort that fires *after* fragments are written. + let moved_addrs = Arc::new(Mutex::new(RoaringTreemap::new())); + let moved_addrs_for_stream = moved_addrs.clone(); + let out_arrow_for_stream = out_arrow.clone(); + // `.then` processes one batch at a time (fully bounded); a small `.buffered(k)` could later + // pipeline the take of the next batch with the write of the current one. + let overlaid = source.then(move |batch| { + let dataset = dataset.clone(); + let plan = plan.clone(); + let moved_addrs = moved_addrs_for_stream.clone(); + async move { + let batch = batch?; + let addresses: Vec = batch[ROW_ADDR] + .as_primitive::() + .values() + .to_vec(); + let full_rows = + TakeBuilder::try_new_from_addresses(dataset, addresses.clone(), plan)? + .execute() + .await?; + let mut overlaid = full_rows; + for field in batch.schema().fields() { + let name = field.name(); + if name == ROW_ADDR || name == ROW_ID { + continue; + } + let new_col = batch + .column_by_name(name) + .expect("updated column present in merger output") + .clone(); + overlaid = overlaid.replace_column_by_name(name, new_col)?; + } + // Only after the take succeeded: a batch that never reaches the writer must not + // tombstone its originals. + moved_addrs + .lock() + .expect("covered move address set poisoned") + .extend(addresses); + Ok::<_, DataFusionError>(overlaid) + } + }); + let write_stream = Box::pin(RecordBatchStreamAdapter::new( + out_arrow_for_stream, + overlaid, + )); + + // Write the moved rows as new fragments. + // Retain the target-base info: `write_fragments_internal` consumes it, but a later cleanup + // must still resolve fragments routed to external bases (otherwise `cleanup_data_fragments` + // skips every file with a `base_id`, orphaning it). + let bases_for_cleanup = target_bases_info.clone(); + let (new_fragments, _) = write_fragments_internal( + self.dataset + .manifest() + .data_storage_format + .lance_file_format(), + Some(&self.dataset), + self.dataset.object_store.clone(), + &self.dataset.base, + full_schema.clone(), + write_stream, + WriteParams::default(), + target_bases_info, + ) + .await?; + + // The overlay recorded every address it handed to the writer; the write above has now + // drained the stream, so the set is complete. + let removed_row_addrs = std::mem::take( + &mut *moved_addrs + .lock() + .expect("covered move address set poisoned"), + ); + if removed_row_addrs.is_empty() { + // Nothing matched -> nothing written; return a no-op update. + return Ok((empty_op(preserving), RoaringTreemap::new())); + } + + // Tombstone the originals. + let deletions = Self::apply_deletions(&self.dataset, &removed_row_addrs).await; + let (updated_fragments, removed_fragment_ids) = match deletions { + Ok(v) => v, + Err(e) => { + cleanup_data_fragments( + &self.dataset.object_store, + &self.dataset.base, + bases_for_cleanup.as_deref(), + &new_fragments, + ) + .await; + return Err(e); + } + }; + + Ok(( + Operation::Update { + removed_fragment_ids, + updated_fragments, + new_fragments, + fields_modified: Vec::new(), + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: preserving, + update_mode: Some(RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }, + removed_row_addrs, + )) + } + /// Patches the columns carried by `source` into the fragments that hold the /// rows it names, and writes the rows with no target address as new /// fragments. @@ -2736,7 +3068,7 @@ impl MergeInsertJob { let joined = self.create_joined_stream(source).await?; let merger = Merger::try_new( self.params.clone(), - source_schema, + source_schema.clone(), !is_full_schema, self.dataset.manifest.uses_stable_row_ids(), )?; @@ -2818,38 +3150,108 @@ impl MergeInsertJob { return Err(Error::not_supported_source("Deleting rows from the target table when there is no match in the source table is not supported when the source data has a different schema than the target data".into())); } - // We will have a different commit path here too, as we are modifying - // fragments rather than writing new ones - let PatchedFragments { - updated_fragments, - new_fragments, - fields_modified, - matched_offsets, - } = Self::update_fragments( - self.dataset.clone(), - Box::pin(stream), - self.dataset.manifest.version + 1, - target_bases_info, - ) - .await?; - - let operation = Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments, - new_fragments, - fields_modified, - compacted_sstables: self.params.compacted_sstables.clone(), - fields_for_preserving_frag_bitmap: vec![], // in-place update do not affect preserving frag bitmap - update_mode: Some(RewriteColumns), - inserted_rows_filter: None, // not implemented for v1 - // The version stamped above is a guess; carry the patched offsets - // so `build_manifest` can re-stamp them at the real commit - // version after a rebase. - updated_fragment_offsets: Some(matched_offsets), + // A covered column's values are materialized in index storage, so patching one in + // place leaves that copy behind. Correctness is NOT what the row-move buys, and the + // comment here used to claim otherwise: `covering_fields` is stored as a suffix of + // `IndexMetadata::fields`, so `Transaction::index_dependent_leaf_ids` already counts + // a covered column as index-dependent and `prune_updated_fields_from_indices` drops + // the whole updated fragment from every index covering it. Disabling the move + // entirely leaves every covered correctness test in this module passing, both + // optimize variants included -- the stale copy is already unreachable, and no + // append-optimize failure mode could be reproduced. + // + // What the move buys is CONFINEMENT. That prune is all-or-nothing per fragment: one + // partial update evicts the ENTIRE fragment from the vector index, including rows the + // merge never matched, and the vectors are the expensive part to rebuild. The move + // tombstones only the matched rows and re-appends them, so the fragment keeps its + // coverage for everything else. Measured on the fixture in + // `test_covered_update_is_fresh_and_confines_index_invalidation`: with the move the + // index bitmap stays `{0, 1}` over fragments `[0, 1, 2]`; patched in place it + // collapses to `{0}`. A recall/cost property rather than a correctness one, then -- + // and that test is the only thing pinning it. + // + // It also makes the two merge_insert paths AGREE, which is the stronger argument. A + // partial source only reaches this branch when every join key has a scalar index + // (`can_use_create_plan` routes the rest to v2), and v2's `create_plan` fills every + // missing dataset column from the target side and writes FULL rows -- already a + // row-move, already confining. Without this branch the same merge would invalidate a + // whole fragment or only the matched rows depending on whether an unrelated scalar + // index happened to exist on the join key. + // + // Known cost, deliberately not addressed here: there is no match-ratio heuristic, so + // when the matched rows are (nearly) a whole fragment the move rewrites every column + // including the vectors for zero confinement gain, where the in-place patch writes one + // column file for the same coverage outcome. Adding one is not cheap -- the strategy + // must be chosen before the stream is consumed, and the per-fragment match + // distribution is not known until the merger has run, so it would take either a + // pre-pass scan or buffering the whole merge, and the latter breaks the bounded-memory + // property this path is built around. + let covered = self.covered_columns_updated(source_schema.as_ref()).await?; + // The move cannot handle every merge (see `covered_move_blocker`). Because it is an + // optimisation over an already-correct in-place patch, a case it cannot handle + // degrades to that patch instead of failing the merge -- otherwise adding + // `covering_columns` to a vector index would silently delete merge_insert shapes that + // worked before, and only on the half of the routing that reaches this code at all. + let move_blocked = if covered.is_empty() { + None + } else { + self.covered_move_blocker(source_schema.as_ref()) }; - // We have rewritten the fragments, not just the deletion files, so - // we can't use affected rows here. - (operation, None) + if !covered.is_empty() && move_blocked.is_none() { + let (operation, removed_row_addrs) = self + .rewrite_covered_update_as_move( + Box::pin(stream) as SendableRecordBatchStream, + target_bases_info, + ) + .await?; + // Same shape as the full-schema `RewriteRows` paths below -- deletions on the + // originals plus new fragments -- so publish the moved addresses. Without them + // `check_update_txn` turns any concurrent delete/update touching these fragments + // into a hard conflict, costing covered datasets concurrency that non-covered + // ones keep. + (operation, Some(RowAddrTreeMap::from(removed_row_addrs))) + } else { + if let Some(reason) = &move_blocked { + info!( + "merge_insert is updating covered column(s) {covered:?} in place rather \ + than moving the matched rows, because {reason}. The updated fragment(s) \ + lose their coverage in every index covering those columns until the next \ + optimize; the merge itself is unaffected." + ); + } + // We will have a different commit path here too, as we are modifying + // fragments rather than writing new ones + let PatchedFragments { + updated_fragments, + new_fragments, + fields_modified, + matched_offsets, + } = Self::update_fragments( + self.dataset.clone(), + Box::pin(stream), + self.dataset.manifest.version + 1, + target_bases_info, + ) + .await?; + + let operation = Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments, + new_fragments, + fields_modified, + compacted_sstables: self.params.compacted_sstables.clone(), + fields_for_preserving_frag_bitmap: vec![], // in-place update do not affect preserving frag bitmap + update_mode: Some(RewriteColumns), + inserted_rows_filter: None, // not implemented for v1 + // The version stamped above is a guess; carry the patched offsets + // so `build_manifest` can re-stamp them at the real commit + // version after a rebase. + updated_fragment_offsets: Some(matched_offsets), + }; + // We have rewritten the fragments, not just the deletion files, so + // we can't use affected rows here. + (operation, None) + } } else { let cleanup_bases = target_bases_info.clone(); let (mut new_fragments, _) = write_fragments_internal( @@ -13248,7 +13650,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n mod external_error { use super::*; - use arrow_schema::{ArrowError, Field as ArrowField, Schema as ArrowSchema}; + use arrow_schema::{ArrowError, Field as ArrowField, Schema}; use std::fmt; #[derive(Debug)] @@ -13267,7 +13669,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_execute_reader_preserves_error_message() { - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ ArrowField::new("key", DataType::Int32, false), ArrowField::new("value", DataType::Int32, false), ])); @@ -13699,16 +14101,14 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n assert_eq!(combined.num_rows(), 300); } - // Regression test: after a partial-schema merge_insert drops a fragment from the vector - // index bitmap, a vector search should not return duplicate rows. The stale vector index - // data still references the dropped fragment, and the scanner also flat-scans unindexed - // fragments, causing the same rows to appear from both paths. - #[tokio::test] - async fn test_partial_merge_insert_stale_vector_index_duplicates() { + /// Build a dataset with a plain (non-covered) IVF_FLAT vector index whose fragment 1 + /// has been pruned from the index bitmap by a partial-schema merge_insert that patched + /// the indexed vector column in place -- the pruned fragment's stale rows still live in + /// the old index segment's storage. Returns `(dataset, total_rows, rows_per_frag, dim)`. + async fn build_pruned_plain_vector_dataset(uri: &str) -> (Arc, usize, usize, i32) { let dim = 4i32; let rows_per_frag = 10usize; let num_frags = 3usize; - let total_rows = rows_per_frag * num_frags; let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), @@ -13745,9 +14145,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n // Write 3 fragments let batch0 = make_batch(0, 0.0); let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); - let mut ds = Dataset::write(reader, "memory://vector_stale_test", None) - .await - .unwrap(); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); for frag_idx in 1..num_frags { let batch = make_batch(frag_idx, 0.0); let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); @@ -13797,7 +14195,18 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .await .unwrap(); - // KNN search with k = total_rows to retrieve all rows + (ds, rows_per_frag * num_frags, rows_per_frag, dim) + } + + /// A KNN query over every row must return each id exactly once -- no stale duplicate + /// served from old index storage alongside the fresh copy. + async fn assert_knn_no_duplicate_ids( + ds: &Dataset, + total_rows: usize, + rows_per_frag: usize, + dim: i32, + ) { + let frag1_start = rows_per_frag; let query: Float32Array = (0..dim) .map(|i| (frag1_start * dim as usize + i as usize) as f32 + 0.5) .collect(); @@ -13809,7 +14218,6 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .await .unwrap(); - // Check no duplicate ids let ids = results .column_by_name("id") .unwrap() @@ -13827,85 +14235,191 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n ); } - // Regression test: after a partial-schema merge_insert drops a fragment from the FTS - // index bitmap, a full text search should not return duplicate rows. The stale inverted - // index data still references the dropped fragment, and the scanner also flat-scans - // unindexed fragments, causing the same rows to appear from both paths. + // Regression test: after a partial-schema merge_insert drops a fragment from the vector + // index bitmap, a vector search should not return duplicate rows. The stale vector index + // data still references the dropped fragment, and the scanner also flat-scans unindexed + // fragments, causing the same rows to appear from both paths. #[tokio::test] - async fn test_partial_merge_insert_stale_fts_index_duplicates() { - let rows_per_frag = 10usize; - let num_frags = 3usize; + async fn test_partial_merge_insert_stale_vector_index_duplicates() { + let (ds, total_rows, rows_per_frag, dim) = + build_pruned_plain_vector_dataset("memory://vector_stale_test").await; + assert_knn_no_duplicate_ids(&ds, total_rows, rows_per_frag, dim).await; + } + + /// Same pruned-fragment state, but merged back into the index: + /// `optimize_indices(merge)` re-scans the pruned fragment as unindexed data AND + /// carries the fragment's stale rows from the old segment's storage -- the merge + /// must drop the stale copies for plain (non-covered) indexes just as it does for + /// covered ones, or ANN queries return duplicate row ids from the merged index. + #[tokio::test] + async fn test_optimize_merge_after_partial_update_no_duplicates_non_covered() { + use lance_index::optimize::OptimizeOptions; + let (ds, total_rows, rows_per_frag, dim) = + build_pruned_plain_vector_dataset("memory://vector_stale_opt_merge").await; + let mut ds = Arc::try_unwrap(ds).unwrap_or_else(|arc| (*arc).clone()); + ds.optimize_indices(&OptimizeOptions::merge(1)) + .await + .unwrap(); + assert_knn_no_duplicate_ids(&ds, total_rows, rows_per_frag, dim).await; + } + + /// A partial merge_insert that rewrites only a covered column is a row-MOVE: the + /// touched rows are tombstoned at their old addresses and reinserted into a new + /// fragment, so the index's stale copy is hidden by the deletion mask rather than + /// patched in place. Fragment 1 is updated in full, so it disappears entirely, and a + /// covered query must observe the fresh values for every row. + #[tokio::test] + async fn test_partial_merge_insert_moves_covered_column_rows() { + let (ds, total_rows, rows_per_frag, dim) = + build_moved_covering_dataset("memory://covering_move_test").await; + assert!( + !ds.get_fragments().iter().any(|f| f.id() == 1), + "fragment 1's rows should have been moved out (tombstoned + reinserted), \ + but fragment 1 is still present" + ); + + assert_covered_query_no_duplicates_fresh(&ds, total_rows, rows_per_frag, dim).await; + } + + /// Pins the two things a partial merge_insert on a covered column must deliver: covered reads + /// are FRESH with no intervening optimize, and the index invalidation is confined to the rows + /// actually updated. + /// + /// The second half is what the row-move uniquely buys, and it is the half that discriminates. + /// Freshness alone does not: `covering_fields` is stored as a suffix of `IndexMetadata::fields`, + /// so `prune_updated_fields_from_indices` already counts a covered column as index-dependent + /// and drops the whole updated fragment from the index whenever a `RewriteColumns` update + /// touches one. That backstop hides the stale copy on its own -- verified by disabling the + /// row-move entirely, which leaves every covered correctness assertion in this module passing. + /// What it costs is coverage: an in-place update of one payload column evicts the ENTIRE + /// fragment from the vector index, including rows the merge never touched, and the vectors are + /// the expensive part to rebuild. The row-move instead tombstones only the matched rows and + /// re-appends them, so the fragment keeps its coverage for everything else. Measured here: + /// with the move the bitmap stays `{0, 1}` over fragments `[0, 1, 2]`; patched in place it + /// collapses to `{0}`. + /// + /// Freshness is still checked, against an INDEPENDENT `take_rows` from the base table rather + /// than against a formula this test also wrote -- asserting that a column came back, or that + /// the query returned something, passes just as happily against a stale copy. `take_rows` + /// re-orders its output to match the requested row ids (see the remapping branch in + /// `dataset::take`), so the two reads line up row for row. + /// + /// Shape follows the recipe in `dataset::tests::dataset_index`: four well-separated clusters, + /// since a covered-ANN test with one partition never reaches the probe path at all and passes + /// against a broken covered read. + #[tokio::test] + async fn test_covered_update_is_fresh_and_confines_index_invalidation() { + use arrow_array::cast::AsArray; + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + use std::collections::HashSet; + + const DIMS: i32 = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + // Only the last cluster is rewritten, so one query spans both moved rows and untouched + // rows still served from the original index storage. + const UPDATED_FROM: i32 = ((NUM_CLUSTERS - 1) * ROWS_PER_CLUSTER) as i32; let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("category", DataType::Utf8, false), - Field::new("text", DataType::Utf8, false), + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIMS, + ), + false, + ), ])); - let make_batch = |frag_idx: usize| { - let start = frag_idx * rows_per_frag; - let ids: Vec = (start..start + rows_per_frag) - .map(|j| format!("id-{j:04}")) - .collect(); - let cats: Vec<&str> = vec!["A"; rows_per_frag]; - // Every row contains "common" so we can search for it and expect all rows - let texts: Vec = (start..start + rows_per_frag) - .map(|j| format!("common unique{j:04}")) - .collect(); + // Cluster centers 1000 apart on dim 0 with deterministic jitter: uniform-random vectors + // leave the early-pruning heuristic searching every partition. + let make_batch = |clusters: std::ops::Range| { + let mut ids = Vec::new(); + let mut payloads = Vec::new(); + let mut values = Vec::new(); + for cluster in clusters { + let center = (cluster * 1000) as f32; + for row in 0..ROWS_PER_CLUSTER { + let n = cluster * ROWS_PER_CLUSTER + row; + ids.push(n as i32); + payloads.push(n as i32 * 10); + for dim in 0..DIMS as usize { + let base = if dim == 0 { center } else { 0.0 }; + values.push(base + ((n * 7919 + dim) % 97) as f32 * 0.0001); + } + } + } + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIMS).unwrap(); RecordBatch::try_new( schema.clone(), vec![ - Arc::new(StringArray::from(ids)), - Arc::new(StringArray::from(cats)), - Arc::new(StringArray::from(texts)), + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(payloads)), + Arc::new(vectors), ], ) .unwrap() }; - // Write 3 fragments - let batch0 = make_batch(0); - let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); - let mut ds = Dataset::write(reader, "memory://fts_stale_test", None) + // Two fragments, so the move has to reach across more than one. + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0..2))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_fresh_no_optimize", None) .await .unwrap(); - for frag_idx in 1..num_frags { - let batch = make_batch(frag_idx); - let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); - ds.append(reader, None).await.unwrap(); - } + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(2..NUM_CLUSTERS))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); - // Create inverted index on text - let params = InvertedIndexParams::default(); - ds.create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string()]); + ds.create_index(&["vector"], IndexType::Vector, None, ¶ms, false) .await .unwrap(); - let ds = Arc::new(ds); - // Partial merge_insert with (id, text) on fragment 1 rows. - // Text still contains "common" so FTS will find them via both paths. - // This drops fragment 1 from the inverted index bitmap. - let frag1_start = rows_per_frag; - let ids: Vec = (frag1_start..frag1_start + rows_per_frag) - .map(|j| format!("id-{j:04}")) - .collect(); - let texts: Vec = (frag1_start..frag1_start + rows_per_frag) - .map(|j| format!("common updated{j:04}")) - .collect(); + // Precondition the assertion depends on: the index really does carry `payload`, so a + // projection of just `payload` is answered from index storage. + let payload_field_id = ds.schema().field("payload").unwrap().id; + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vector_idx").unwrap(); + assert_eq!(vec_idx.covering_fields, vec![payload_field_id]); + + // Partial-schema merge_insert rewriting ONLY the covered column, for the last cluster. let sub_schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("text", DataType::Utf8, false), + Field::new("id", DataType::Int32, false), + Field::new("payload", DataType::Int32, false), ])); + let updated_ids: Vec = (UPDATED_FROM..TOTAL as i32).collect(); + let updated_payloads: Vec = updated_ids.iter().map(|id| id * 10 + 7).collect(); let update_batch = RecordBatch::try_new( sub_schema.clone(), vec![ - Arc::new(StringArray::from(ids)), - Arc::new(StringArray::from(texts)), + Arc::new(Int32Array::from(updated_ids)), + Arc::new(Int32Array::from(updated_payloads)), ], ) .unwrap(); let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let frags_before: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) .unwrap() .when_matched(WhenMatched::UpdateAll) @@ -13915,17 +14429,1653 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n .execute_reader(reader) .await .unwrap(); + // Deliberately no `optimize_indices` here: the point is that the stale copy is hidden + // without one. - // FTS search for "common" — every row should match exactly once - let query = FullTextSearchQuery::new("common".to_string()); - let results = ds + // The move half: the matched rows left fragment 1 for a brand-new fragment, and + // fragment 1 -- which also holds 64 rows the merge never matched -- keeps its place in + // the vector index. An in-place patch evicts it instead. + let frags_after: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + assert!( + frags_after.difference(&frags_before).count() == 1, + "the matched rows should have been re-appended to one new fragment; \ + fragments went from {frags_before:?} to {frags_after:?}" + ); + let updated_frag = *frags_before.iter().max().unwrap(); + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vector_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + bitmap.contains(updated_frag as u32), + "a covered update must not evict the whole fragment from the vector index: \ + fragment {updated_frag} still holds {ROWS_PER_CLUSTER} untouched rows, but the \ + index now covers only {bitmap:?}" + ); + + let mut q_values = vec![0.0f32; DIMS as usize]; + q_values[0] = ((NUM_CLUSTERS - 1) * 1000) as f32; + let query = Float32Array::from(q_values); + + // Project ONLY the covered column. Adding an uncovered column here would send the whole + // projection to the base table and the covered storage would never be read, so the test + // would pass against a stale copy. + let mut scan = ds.scan(); + scan.nearest("vector", &query, TOTAL).unwrap(); + scan.with_row_id(); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNIvfPartition"), + "the covered values must come from the vector index, not a flat rescan:\n{plan}" + ); + let results = scan.try_into_batch().await.unwrap(); + + assert_eq!( + results.num_rows(), + TOTAL, + "covered query returned {} rows, expected {TOTAL}", + results.num_rows() + ); + let row_ids = results[ROW_ID].as_primitive::(); + let unique: HashSet = row_ids.values().iter().copied().collect(); + assert_eq!( + unique.len(), + TOTAL, + "a stale copy survived at its old address: duplicate row ids in covered results" + ); + + // The independent read: same rows, fetched from the base table by row id. + let queried_row_ids: Vec = row_ids.values().to_vec(); + let base = ds + .take_rows( + &queried_row_ids, + ProjectionRequest::from_schema(ds.schema().project(&["payload"]).unwrap()), + ) + .await + .unwrap(); + assert_eq!(base.num_rows(), TOTAL); + + let covered_payload = results["payload"].as_primitive::(); + let base_payload = base["payload"].as_primitive::(); + let mut updated_seen = 0usize; + for i in 0..TOTAL { + assert_eq!( + covered_payload.value(i), + base_payload.value(i), + "covered read disagrees with the base table for row id {}: index says {}, \ + table says {}", + row_ids.value(i), + covered_payload.value(i), + base_payload.value(i) + ); + if base_payload.value(i) % 10 == 7 { + updated_seen += 1; + } + } + // Guards the assertion above from passing vacuously if the merge silently no-ops. + assert_eq!( + updated_seen, ROWS_PER_CLUSTER, + "expected {ROWS_PER_CLUSTER} rewritten rows in the covered results, saw {updated_seen}" + ); + + // The comparison above would still agree if the move had taken the WRONG addresses -- + // the base table would carry the same corruption, so both reads would match. Row + // integrity is what catches that: the move rebuilds each row from a full-row take at an + // address, so a wrongly addressed take pairs one row's `id` with another's payload and + // duplicates or drops ids. + let full = ds .scan() - .full_text_search(query) + .project(&["id", "payload"]) .unwrap() .try_into_batch() .await .unwrap(); - + assert_eq!(full.num_rows(), TOTAL, "the move changed the row count"); + let full_ids = full["id"].as_primitive::(); + let full_payload = full["payload"].as_primitive::(); + let mut seen_ids = HashSet::new(); + for i in 0..TOTAL { + let id = full_ids.value(i); + assert!(seen_ids.insert(id), "id {id} duplicated by the move"); + let expected = if id >= UPDATED_FROM { + id * 10 + 7 + } else { + id * 10 + }; + assert_eq!( + full_payload.value(i), + expected, + "row {id} came back paired with another row's payload" + ); + } + assert_eq!(seen_ids.len(), TOTAL, "the move dropped rows"); + } + + /// Bounded-streaming regression for the covered row-move: a covered-column update spanning + /// MORE than one scan batch (> `BATCH_SIZE_FALLBACK` = 8192 rows) must move every row exactly + /// once with its fresh value. `rewrite_covered_update_as_move` streams the take/overlay/write + /// per batch and accumulates the tombstoned addresses across batches, so this exercises the + /// multi-batch accumulation path the small covered tests never reach. + #[tokio::test] + async fn test_covered_update_move_streams_multiple_batches() { + use arrow_array::cast::AsArray; + use std::collections::HashSet; + + let dim = 2i32; + let rows_per_frag = 5000usize; // 2 frags => 10_000 rows > BATCH_SIZE_FALLBACK (8192) + let num_frags = 2usize; + let total_rows = rows_per_frag * num_frags; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:05}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:05}")) + .collect(); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap() + }; + + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_move_multibatch", None) + .await + .unwrap(); + for frag_idx in 1..num_frags { + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(frag_idx))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + } + + // BTree on the join key + a vector index covering `category`. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial merge_insert rewriting `category` for EVERY row -> a >1-batch covered move. + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:05}")).collect(); + let new_cats: Vec = (0..total_rows).map(|j| format!("updated-{j:05}")).collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Every row must survive exactly once with its fresh category -- nothing lost or + // duplicated across the batch boundary the streaming move crosses. + let batch = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + batch.num_rows(), + total_rows, + "row count changed after a multi-batch covered move" + ); + let ids_col = batch["id"].as_string::(); + let cats_col = batch["category"].as_string::(); + let mut seen = HashSet::new(); + for i in 0..batch.num_rows() { + let id = ids_col.value(i); + assert!( + seen.insert(id.to_string()), + "duplicate id {id} after covered move" + ); + let n: usize = id[3..].parse().unwrap(); + assert_eq!( + cats_col.value(i), + format!("updated-{n:05}"), + "stale or missing covered value for {id}" + ); + } + assert_eq!(seen.len(), total_rows); + } + + /// Runs a partial-schema merge_insert whose source carries a STRUCT with only one of the + /// target's two children, alongside `category`. `cover` decides whether the vector index + /// includes `category`, i.e. whether the covered row-move branch is taken at all. + async fn partial_struct_source_merge(uri: &str, cover: bool) -> Result> { + const ROWS: usize = 256; + const DIM: i32 = 2; + let struct_fields: arrow_schema::Fields = vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ] + .into(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("s", DataType::Struct(struct_fields.clone()), false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM), + false, + ), + ])); + let ids: Vec = (0..ROWS).map(|j| format!("id-{j:03}")).collect(); + let cats: Vec = (0..ROWS).map(|j| format!("cat-{j:03}")).collect(); + let a = Int32Array::from((0..ROWS as i32).collect::>()); + let b = Int32Array::from((0..ROWS as i32).map(|v| v + 1000).collect::>()); + let s = StructArray::new(struct_fields, vec![Arc::new(a), Arc::new(b)], None); + let values: Vec = (0..ROWS * DIM as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(s), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + if cover { + params.covering_columns(vec!["category".to_string()]); + } + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial source: the join key, the (maybe-covered) `category`, and a STRUCT with + // only one of its two children. + let partial_struct_fields: arrow_schema::Fields = + vec![Field::new("a", DataType::Int32, false)].into(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("s", DataType::Struct(partial_struct_fields.clone()), false), + ])); + let new_cats: Vec = (0..ROWS).map(|j| format!("upd-{j:03}")).collect(); + let new_a = Int32Array::from((0..ROWS as i32).map(|v| v + 500).collect::>()); + let new_s = StructArray::new(partial_struct_fields, vec![Arc::new(new_a)], None); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + Arc::new(new_s), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .map(|(ds, _)| ds) + } + + /// `id -> (s.a, s.b, category)` for the whole dataset, so the covered and uncovered runs of + /// `partial_struct_source_merge` can be compared row for row. + async fn partial_struct_merge_contents(ds: &Dataset) -> HashMap { + let batch = ds + .scan() + .project(&["id", "category", "s"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_string::(); + let cats = batch["category"].as_string::(); + let s = batch["s"].as_struct(); + let a = s.column_by_name("a").unwrap().as_primitive::(); + let b = s.column_by_name("b").unwrap().as_primitive::(); + (0..batch.num_rows()) + .map(|i| { + ( + ids.value(i).to_string(), + (a.value(i), b.value(i), cats.value(i).to_string()), + ) + }) + .collect() + } + + /// The covered row-move overlays each source column onto a full-row take by REPLACING the + /// whole column, so it needs the source column's Arrow type to equal the target's exactly. A + /// partial-schema source is only required to be a *subschema*, and `compare_fields` recurses + /// with the same options, so a struct may legally arrive carrying only some of its children. + /// Left unguarded this fails with an opaque Arrow "column types must match schema types" from + /// `replace_column_by_name`; guarded, it falls back to the in-place patch. + /// + /// The control is the whole point: the same merge on a dataset whose index does NOT cover + /// `category` takes the in-place branch, and both runs must produce **identical data**. That + /// pins the fallback as behaviour-preserving rather than merely non-failing -- adding + /// `covering_columns` to a vector index must not change what a merge does, only how much index + /// coverage survives it. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_with_partial_struct_source() { + let plain = partial_struct_source_merge("memory://partial_struct_plain", false) + .await + .expect( + "control: a partial struct source is handled by the in-place branch, so if this \ + fails the comparison below is no longer evidence of anything", + ); + let covered = partial_struct_source_merge("memory://partial_struct_covered", true) + .await + .expect("a partial struct source must fall back to in-place, not fail"); + + let plain_rows = partial_struct_merge_contents(&plain).await; + let covered_rows = partial_struct_merge_contents(&covered).await; + assert_eq!(plain_rows.len(), 256); + assert_eq!( + plain_rows, covered_rows, + "covering a column must not change what the merge writes" + ); + // Spot-check against the expected values too, so an identical-but-wrong pair of runs + // cannot pass: `s.a` was updated, `s.b` was never in the source and must survive. + assert_eq!(covered_rows["id-007"], (507, 1007, "upd-007".to_string())); + + // Confinement is what covering forfeits here, and only that. + let indices = covered.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + bitmap.is_empty(), + "the in-place fallback should have evicted the only fragment from the vector index; \ + still covered means the fallback was not taken: {bitmap:?}" + ); + } + + /// The covered row-move must not depend on the order its source arrives in: it accumulates + /// tombstoned addresses with `extend` rather than the bulk `RoaringTreemap::append`, which + /// rejects any value <= the running maximum and would abort *after* fragments were written, + /// orphaning them. + /// + /// Note what this can and cannot reach. The BTree is required, not incidental: without a + /// scalar index on every join key, `can_use_create_plan` routes the merge to v2 and this + /// function never runs (the earlier version of this test omitted the index specifically to + /// get the hash join, and so tested nothing here). With it, the merge takes the indexed-scan + /// join -- the only route to the move -- which reads in fragment order, so descending + /// addresses cannot be produced through the public API at all. What is pinned is that a + /// reversed *source* is handled correctly; `extend` remains a deliberate defence against a + /// future routing change rather than something reachable today. + #[tokio::test] + async fn test_covered_update_move_accepts_unsorted_source_order() { + use arrow_array::cast::AsArray; + use std::collections::HashSet; + + let dim = 2i32; + let total_rows = 128usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..total_rows).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..total_rows * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_move_unsorted", None) + .await + .unwrap(); + // The BTree is what makes this test reach `rewrite_covered_update_as_move` at all -- + // see the doc comment. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Feed the source fully reversed. + let mut update_ids: Vec = ids.clone(); + update_ids.reverse(); + let new_cats: Vec = update_ids + .iter() + .map(|id| format!("updated-{}", &id[3..])) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(update_ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let frags_before: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Proof this test reached the move rather than the in-place patch: the move re-appends the + // matched rows to a new fragment, where the in-place patch leaves the fragment ids alone. + // Without this the value checks below pass on either path. + // + // This does NOT by itself exclude the v2 path, which also writes full rows and would also + // append. v2 is excluded structurally instead: the scalar index created above means + // `can_use_create_plan`'s `!would_use_scalar_index` is false. That coupling is load-bearing + // -- it is exactly what this test lacked when it silently never ran the move at all. + let frags_after: HashSet = ds.get_fragments().iter().map(|f| f.id()).collect(); + assert!( + !frags_after + .difference(&frags_before) + .collect::>() + .is_empty(), + "expected the covered row-move to append a new fragment; fragments went from \ + {frags_before:?} to {frags_after:?}, so this test never ran the move" + ); + + let batch = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + batch.num_rows(), + total_rows, + "row count changed after an out-of-order covered move" + ); + let ids_col = batch["id"].as_string::(); + let cats_col = batch["category"].as_string::(); + let mut seen = HashSet::new(); + for i in 0..batch.num_rows() { + let id = ids_col.value(i); + assert!( + seen.insert(id.to_string()), + "duplicate id {id} after covered move" + ); + assert_eq!( + cats_col.value(i), + format!("updated-{}", &id[3..]), + "stale or missing covered value for {id}" + ); + } + assert_eq!(seen.len(), total_rows); + } + + /// The covered row-move must publish the moved rows as `affected_rows`, exactly as the + /// full-schema `RewriteRows` paths do. Without them `check_update_txn` short-circuits on + /// `affected_rows.is_none()` and turns any concurrent delete/update touching the same + /// fragments into a retryable conflict, so covered datasets would silently lose the + /// merge_insert/delete concurrency that non-covered ones keep. + #[tokio::test] + async fn test_covered_update_move_publishes_affected_rows() { + let dim = 2i32; + let total_rows = 64usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..total_rows).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..total_rows).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..total_rows * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_move_affected_rows", None) + .await + .unwrap(); + // Required to reach `rewrite_covered_update_as_move`: without a scalar index on every + // join key, `can_use_create_plan` routes the merge to v2, which publishes affected rows + // of its own and would satisfy this test's assertion without the move running at all. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Single fragment written in id order, so row N sits at address N. + let targets = [2usize, 5, 7]; + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from( + targets + .iter() + .map(|n| format!("id-{n:04}")) + .collect::>(), + )), + Arc::new(StringArray::from( + targets + .iter() + .map(|n| format!("updated-{n:04}")) + .collect::>(), + )), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let UncommittedMergeInsert { affected_rows, .. } = + MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_uncommitted(reader) + .await + .unwrap(); + + let expected = + RowAddrTreeMap::from(RoaringTreemap::from_iter(targets.iter().map(|n| *n as u64))); + assert_eq!( + affected_rows, + Some(expected), + "covered row-move must publish the moved addresses so a concurrent delete can rebase" + ); + } + + /// Builds a 3-fragment dataset with an IVF_PQ index covering `category`, then does a + /// partial merge_insert that rewrites only `category` for fragment 1. That update is a + /// row-MOVE: fragment 1's rows are tombstoned and rewritten into a new fragment, while + /// the old segment's storage still physically holds their pre-update covering values. + /// Returns `(post-move dataset, total_rows, rows_per_frag, dim)`. + /// The covered move must report EVERY field id -- leaves included -- in + /// `fields_for_preserving_frag_bitmap`. The consumer compares that list against + /// each index's leaf-expanded dependency set, so a top-level-only list would let + /// an index keyed on a nested field (its dependency set holds the LEAF id) + /// escape the exclude-everything contract and get its bitmap wrongly extended + /// onto the moved fragment. Unobservable through query results today (the move + /// runs only on address-style datasets, where the consumer skips every index), + /// so this pins the committed transaction payload itself. + #[tokio::test] + async fn test_covered_move_reports_leaf_field_ids_for_preserving() { + use crate::io::commit::read_transaction_file; + + let dim = 4i32; + let rows_per_frag = 256usize; // >= 256 rows so PQ has enough to train + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "meta", + DataType::Struct(vec![Field::new("code", DataType::Int32, false)].into()), + false, + ), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:04}")) + .collect(); + let codes = + Int32Array::from_iter_values((start..start + rows_per_frag).map(|j| j as i32)); + let meta = StructArray::from(vec![( + Arc::new(Field::new("code", DataType::Int32, false)), + Arc::new(codes) as arrow_array::ArrayRef, + )]); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(meta), + Arc::new(vectors), + ], + ) + .unwrap() + }; + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, "memory://covered_move_preserving", None) + .await + .unwrap(); + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(1))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + // Scalar index on the join key so the merge joins via the indexed scan and + // reaches `rewrite_covered_update_as_move` (the full-scan join commits its + // move through the insert exec instead, which is not the producer under + // test here). + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial covered-column update over fragment 1 -> takes the row-move path. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let new_cats: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("updated-{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + let tx_path = ds + .manifest() + .transaction_file + .clone() + .expect("the covered move writes a transaction file"); + let tx = read_transaction_file(ds.object_store.as_ref(), &ds.base, &tx_path) + .await + .unwrap(); + let Operation::Update { + fields_for_preserving_frag_bitmap, + update_mode, + .. + } = &tx.operation + else { + panic!("expected Operation::Update, got: {:?}", tx.operation); + }; + assert_eq!(*update_mode, Some(RewriteRows), "the move path committed"); + let reported: HashSet = fields_for_preserving_frag_bitmap.iter().copied().collect(); + let expected: HashSet = ds + .schema() + .fields_pre_order() + .map(|f| f.id as u32) + .collect(); + let leaf_id = ds.schema().field("meta").unwrap().children[0].id as u32; + assert!( + reported.contains(&leaf_id), + "the preserving list must include nested LEAF ids (missing {leaf_id}): {reported:?}" + ); + assert_eq!( + reported, expected, + "the covered move must report every field id, leaves included" + ); + } + + async fn build_moved_covering_dataset(uri: &str) -> (Arc, usize, usize, i32) { + let dim = 4i32; + let rows_per_frag = 256usize; + let num_frags = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec = (start..start + rows_per_frag) + .map(|j| format!("cat-{j:04}")) + .collect(); + let values: Vec = (0..rows_per_frag * dim as usize) + .map(|i| (start * dim as usize + i) as f32) + .collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap() + }; + + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(0))], + schema.clone(), + )); + let mut ds = Dataset::write(reader, uri, None).await.unwrap(); + for frag_idx in 1..num_frags { + let reader = Box::new(RecordBatchIterator::new( + [Ok(make_batch(frag_idx))], + schema.clone(), + )); + ds.append(reader, None).await.unwrap(); + } + + // Scalar index on the join key so the merge joins via the indexed scan. The write + // strategy is unaffected: a covered-column update is always a row-move. + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + let ds = Arc::new(ds); + + // Partial merge_insert rewriting ONLY `category` for fragment 1 -> prunes + // fragment 1 from the vector index's coverage bitmap. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let new_cats: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("updated-{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // Precondition every caller relies on: the covering column really is recorded on + // the index metadata, since the move/prune logic keys on it. + let category_field_id = ds.schema().field("category").unwrap().id; + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + assert_eq!(vec_idx.covering_fields, vec![category_field_id]); + + (ds, num_frags * rows_per_frag, rows_per_frag, dim) + } + + /// A covered query over every row must return each row exactly once (no stale duplicate + /// from the old storage) with the fresh `category` value. + async fn assert_covered_query_no_duplicates_fresh( + ds: &Dataset, + total_rows: usize, + rows_per_frag: usize, + dim: i32, + ) { + use arrow_array::cast::AsArray; + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + use std::collections::HashSet; + + let frag1_start = rows_per_frag; + let query: Float32Array = (0..dim) + .map(|i| (frag1_start * dim as usize + i as usize) as f32) + .collect(); + let results = ds + .scan() + .nearest("vec", &query, total_rows) + .unwrap() + .with_row_id() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + let rowids = results[ROW_ID].as_primitive::(); + let unique: HashSet = rowids.values().iter().copied().collect(); + assert_eq!( + unique.len(), + rowids.len(), + "optimize double-counted the pruned fragment: duplicate row ids in covered results" + ); + assert_eq!( + results.num_rows(), + total_rows, + "covered query returned {} rows, expected {total_rows} (dropped/duplicated rows)", + results.num_rows() + ); + + let ids_col = results["id"].as_string::(); + let cats = results["category"].as_string::(); + for i in 0..ids_col.len() { + let id = ids_col.value(i); + let n: usize = id[3..].parse().unwrap(); + let expected = if (frag1_start..frag1_start + rows_per_frag).contains(&n) { + format!("updated-{n:04}") + } else { + format!("cat-{n:04}") + }; + assert_eq!(cats.value(i), expected, "stale covering value for row {id}"); + } + } + + /// After a covered-column update moves rows out of a fragment, an optimize must not + /// resurrect the stale copies still sitting in the old segment's storage. `merge` + /// re-scans the moved rows as unindexed data *and* folds in the old segment, so it must + /// not double-count them; the default (append) optimize leaves the old segment in place, + /// so the deletion mask must keep hiding them. + #[rstest::rstest] + #[case::merge(lance_index::optimize::OptimizeOptions::merge(1))] + #[case::append_default(lance_index::optimize::OptimizeOptions::default())] + #[tokio::test] + async fn test_optimize_after_covering_move_no_duplicate_rows( + #[case] options: lance_index::optimize::OptimizeOptions, + ) { + let (ds, total_rows, rows_per_frag, dim) = + build_moved_covering_dataset("memory://covered_opt").await; + let mut ds = Arc::try_unwrap(ds).unwrap_or_else(|arc| (*arc).clone()); + ds.optimize_indices(&options).await.unwrap(); + assert_covered_query_no_duplicates_fresh(&ds, total_rows, rows_per_frag, dim).await; + } + + /// The covered move cannot carry stable row ids onto its new fragments -- it never rechunks + /// the captured `row_id_sequence` the way the full-schema path does, so `assign_row_ids` + /// would mint fresh ids for every moved row and break the one guarantee stable row ids exist + /// to provide. The move is an optimisation over an already-correct in-place patch, so this + /// must degrade to that patch rather than fail the merge. + /// + /// The assertions are the property the old rejection protected, stated positively: the merge + /// SUCCEEDS, every row keeps the row id it had (the in-place patch moves no rows, so ids are + /// preserved by construction -- strictly safer here than the move), and confinement was the + /// only thing given up, i.e. the updated fragment left the vector index's bitmap. Run against + /// a real directory rather than `memory://` so the ids are checked after a genuine reopen. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_on_stable_row_ids() { + use arrow_array::types::UInt64Type; + use lance_core::ROW_ID; + + /// `id -> _rowid` for the whole dataset, so "ids were preserved" is provable per row + /// rather than as a count. + async fn row_ids_by_key(ds: &Dataset) -> HashMap { + let batch = ds + .scan() + .with_row_id() + .project(&["id"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = batch["id"].as_string::(); + let row_ids = batch[ROW_ID].as_primitive::(); + (0..batch.num_rows()) + .map(|i| (ids.value(i).to_string(), row_ids.value(i))) + .collect() + } + + let test_uri = TempStrDir::default(); + let dim = 4i32; + let n = 256usize; // enough rows to train PQ + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_pq(1, 8, 2, MetricType::L2, 2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial merge_insert updating only the covered `category`. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from( + (0..n).map(|j| format!("new-{j:04}")).collect::>(), + )), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let row_ids_before = row_ids_by_key(&ds).await; + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update on a stable-row-id dataset must fall back, not fail"); + + // Reopen so the ids are read back from committed bytes, not from the in-memory handle. + let reopened = Dataset::open(test_uri.as_ref()).await.unwrap(); + let after = reopened + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(after.num_rows(), n, "the fallback must not lose rows"); + let cats = after["category"].as_string::(); + assert!( + (0..cats.len()).all(|i| cats.value(i).starts_with("new-")), + "the fallback must still apply the update" + ); + + // The property the old rejection existed to protect: no row changed identity. The + // in-place patch moves no rows, so every id survives, paired with the same key. + let row_ids_after = row_ids_by_key(&reopened).await; + assert_eq!( + row_ids_before, row_ids_after, + "falling back to the in-place patch must preserve every stable row id" + ); + + // And the only thing forfeited is confinement: the updated fragment left the bitmap. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; if it is still covered the merge did not take the fallback and this test is \ + not exercising it: {bitmap:?}" + ); + } + + /// A partial-schema merge that does NOT update matched rows (DoNothing + InsertAll) + /// must not be treated as a covered-column update, even when the source carries a + /// covered column -- matched rows are untouched and unmatched rows enter unindexed + /// fragments, so nothing is patched into the index. + #[tokio::test] + async fn test_merge_insert_covered_donothing_insert_not_rejected() { + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + // nullable so InsertAll with a partial source is permitted + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_donothing", None) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // DoNothing on matched + InsertAll unmatched, partial source (id + category) + // including new rows to insert. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let new_ids: Vec = (n..n + 5).map(|j| format!("id-{j:04}")).collect(); + let new_cats: Vec = (n..n + 5).map(|j| format!("cat-{j:04}")).collect(); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(new_ids)), + Arc::new(StringArray::from(new_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let result = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::DoNothing) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(reader) + .await; + assert!( + result.is_ok(), + "DoNothing + InsertAll with a covered column in a partial source must not be \ + rejected as a covered update; got {:?}", + result.err() + ); + } + + /// An upsert touching a covered column -- `UpdateAll` matched plus `InsertAll` unmatched -- + /// must fall back to the in-place patch. The move reads row addresses through the null mask, + /// and the unmatched group arrives with a NULL `_rowaddr`, so entering the move would + /// tombstone unrelated rows at address 0. Never entering it prevents that as completely as + /// refusing the merge did, and `update_fragments` already routes the NULL group to + /// `handle_new_fragments`. + /// + /// The data-loss assertion is the point: the row count must be exactly the original plus the + /// inserted row, and every pre-existing row must still be there. That is what a bogus + /// address-0 tombstone would break. + #[tokio::test] + async fn test_merge_insert_covered_update_with_insert_falls_back() { + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + // nullable so InsertAll with a partial source is permitted + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://covered_update_insert", None) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial source mixing an update to an existing row with a brand-new row. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + "id-0001".to_string(), + format!("id-{:04}", n), + ])), + Arc::new(StringArray::from(vec![ + "updated-0001".to_string(), + "brand-new".to_string(), + ])), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + let (ds, stats) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update carrying inserts must fall back, not fail"); + assert_eq!(stats.num_updated_rows, 1); + assert_eq!(stats.num_inserted_rows, 1); + + // No row was tombstoned at a bogus address: exactly the original rows plus the insert, + // with every original key still present. + let after = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + after.num_rows(), + n + 1, + "the fallback must insert the new row and tombstone nothing else" + ); + let ids_col = after["id"].as_string::(); + let cats_col = after["category"].as_string::(); + let found: HashMap<&str, &str> = (0..after.num_rows()) + .map(|i| (ids_col.value(i), cats_col.value(i))) + .collect(); + for j in 0..n { + assert!( + found.contains_key(format!("id-{j:04}").as_str()), + "the fallback dropped pre-existing row id-{j:04}" + ); + } + assert_eq!( + found["id-0001"], "updated-0001", + "the update was not applied" + ); + assert_eq!( + found[format!("id-{n:04}").as_str()], + "brand-new", + "the insert was not applied" + ); + + // Confinement is what was forfeited. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; still covered means the fallback was not taken: {bitmap:?}" + ); + } + + /// The covered-column row-move re-reads full rows via a take, which cannot round-trip + /// legacy (v1) blob columns as binary, so a dataset carrying one must fall back to the + /// in-place patch -- which never reads the blob column at all. The assertion that matters + /// is therefore that the blob survives the merge intact, not that the merge was refused. + #[tokio::test] + async fn test_merge_insert_covered_update_falls_back_with_v1_blob() { + use arrow_array::LargeBinaryArray; + use lance_arrow::BLOB_META_KEY; + + let dim = 4i32; + let n = 20usize; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( + BLOB_META_KEY.to_string(), + "true".to_string(), + )])), + ])); + let ids: Vec = (0..n).map(|j| format!("id-{j:04}")).collect(); + let cats: Vec = (0..n).map(|j| format!("cat-{j:04}")).collect(); + let blobs: Vec> = (0..n).map(|_| Some(b"payload".as_slice())).collect(); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(vectors), + Arc::new(LargeBinaryArray::from(blobs)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + let mut ds = Dataset::write( + reader, + "memory://covered_v1_blob", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }), + ) + .await + .unwrap(); + ds.create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let mut params = VectorIndexParams::ivf_flat(2, MetricType::L2); + params.covering_columns(vec!["category".to_string()]); + ds.create_index(&["vec"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Partial source updating the covered `category` on matched rows only. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + ])); + let upd_ids: Vec = (0..3).map(|j| format!("id-{j:04}")).collect(); + let upd_cats: Vec = (0..3).map(|j| format!("new-{j:04}")).collect(); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(upd_ids)), + Arc::new(StringArray::from(upd_cats)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update)], sub_schema)); + let updated_frag = ds.get_fragments().iter().map(|f| f.id()).max().unwrap(); + // A v1 blob column scans as a position/size description struct, not binary -- writing + // those descriptors back as data is exactly the corruption the move would cause. Capture + // them so "the blob was left alone" is provable rather than assumed. + let blobs_before = ds + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let (ds, _) = MergeInsertBuilder::try_new(Arc::new(ds), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .expect("a covered update on a v1-blob dataset must fall back, not fail"); + + // The update landed, and the blob the move would have corrupted is untouched. Read the + // blobs through the binary path the v1 column is meant to be read by. + let after = ds + .scan() + .project(&["id", "category"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(after.num_rows(), n); + let ids_col = after["id"].as_string::(); + let cats_col = after["category"].as_string::(); + for i in 0..after.num_rows() { + let key: usize = ids_col.value(i)[3..].parse().unwrap(); + let expected = if key < 3 { + format!("new-{key:04}") + } else { + format!("cat-{key:04}") + }; + assert_eq!(cats_col.value(i), expected, "wrong category after fallback"); + } + let blobs_after = ds + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + blobs_before, blobs_after, + "the in-place fallback must leave the v1 blob column untouched; the move would have \ + rewritten these descriptors into a new fragment as ordinary struct data" + ); + // And the bytes are still readable through the blob API they were written with. + let blob_handles = ds + .take_blobs_by_indices(&(0..n as u64).collect::>(), "blobs") + .await + .unwrap(); + assert_eq!(blob_handles.len(), n); + assert!( + blob_handles.iter().all(|b| b + .as_ref() + .is_some_and(|b| b.size() == b"payload".len() as u64)), + "every v1 blob must still read back at its original size" + ); + + // Confinement is what was forfeited. + let indices = ds.load_indices().await.unwrap(); + let vec_idx = indices.iter().find(|i| i.name == "vec_idx").unwrap(); + let bitmap = vec_idx.fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(updated_frag as u32), + "the in-place fallback should have evicted fragment {updated_frag} from the vector \ + index; still covered means the fallback was not taken: {bitmap:?}" + ); + } + + // Regression test: after a partial-schema merge_insert drops a fragment from the FTS + // index bitmap, a full text search should not return duplicate rows. The stale inverted + // index data still references the dropped fragment, and the scanner also flat-scans + // unindexed fragments, causing the same rows to appear from both paths. + #[tokio::test] + async fn test_partial_merge_insert_stale_fts_index_duplicates() { + let rows_per_frag = 10usize; + let num_frags = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + + let make_batch = |frag_idx: usize| { + let start = frag_idx * rows_per_frag; + let ids: Vec = (start..start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let cats: Vec<&str> = vec!["A"; rows_per_frag]; + // Every row contains "common" so we can search for it and expect all rows + let texts: Vec = (start..start + rows_per_frag) + .map(|j| format!("common unique{j:04}")) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(cats)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap() + }; + + // Write 3 fragments + let batch0 = make_batch(0); + let reader = Box::new(RecordBatchIterator::new([Ok(batch0)], schema.clone())); + let mut ds = Dataset::write(reader, "memory://fts_stale_test", None) + .await + .unwrap(); + for frag_idx in 1..num_frags { + let batch = make_batch(frag_idx); + let reader = Box::new(RecordBatchIterator::new([Ok(batch)], schema.clone())); + ds.append(reader, None).await.unwrap(); + } + + // Create inverted index on text + let params = InvertedIndexParams::default(); + ds.create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let ds = Arc::new(ds); + + // Partial merge_insert with (id, text) on fragment 1 rows. + // Text still contains "common" so FTS will find them via both paths. + // This drops fragment 1 from the inverted index bitmap. + let frag1_start = rows_per_frag; + let ids: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("id-{j:04}")) + .collect(); + let texts: Vec = (frag1_start..frag1_start + rows_per_frag) + .map(|j| format!("common updated{j:04}")) + .collect(); + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + let update_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(StringArray::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + let reader = Box::new(RecordBatchIterator::new([Ok(update_batch)], sub_schema)); + let (ds, _) = MergeInsertBuilder::try_new(ds, vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(reader) + .await + .unwrap(); + + // FTS search for "common" — every row should match exactly once + let query = FullTextSearchQuery::new("common".to_string()); + let results = ds + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + // Check no duplicate ids let ids = results .column_by_name("id") @@ -14921,11 +17071,11 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_with_blob_v1_source_provides_blob() { use arrow_array::LargeBinaryArray; - use arrow_schema::Schema as ArrowSchema; + use arrow_schema::Schema; use lance_arrow::BLOB_META_KEY; let test_dir = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ Field::new("blobs", DataType::LargeBinary, true).with_metadata(HashMap::from([( BLOB_META_KEY.to_string(), "true".to_string(), @@ -15000,10 +17150,10 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n #[tokio::test] async fn test_merge_insert_with_blob_v2_source_provides_blob() { use crate::{BlobArrayBuilder, blob_field}; - use arrow_schema::Schema as ArrowSchema; + use arrow_schema::Schema; let test_dir = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ + let schema = Arc::new(Schema::new(vec![ blob_field("blobs", true), Field::new("id", DataType::Int64, true), Field::new("other", DataType::Int64, true), diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 6cc7affee02..ada522328bb 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -137,7 +137,7 @@ fn validate_segment_metadata(index_name: &str, segments: &[IndexMetadata]) -> Re Ok(()) } -fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { +pub(crate) fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { field_ids.insert(field.id); for child in &field.children { collect_subtree_field_ids(child, field_ids); @@ -273,7 +273,32 @@ async fn prune_stale_segment_coverage( .iter_mut() .filter(|segment| segment.dataset_version() == version) { - let indexed_field_ids = segment_indexed_field_ids(dataset, segment)?; + // `segment_indexed_field_ids` already walks every entry of `segment.fields()`, + // keyed and carried alike -- `covering_fields` is always the trailing slice of + // `fields` -- so no separate expansion over `covering_fields()` is needed here. + let relevant_field_ids = segment_indexed_field_ids(dataset, segment)?; + for included_id in segment.covering_fields() { + // A covered field whose subtree changed between the segment's build version + // and now -- e.g. a struct that gained a child via `add_columns` (which + // commits while this segment is still uncommitted, so the Merge-commit guard + // never saw it) -- makes the segment's stored covering payload type- + // incompatible with the current schema for *every* row it covers. No + // per-fragment file-path prune below can repair a global type change, so + // reject the commit; the segment must be rebuilt against the current schema. + if Transaction::covered_field_subtree_changed( + historical.schema(), + dataset.schema(), + *included_id, + ) { + return Err(Error::invalid_input(format!( + "CreateIndex: covered (included) field id {included_id} changed its \ + subtree between the segment's build version {version} and the current \ + dataset version {}; the segment's stored covering payload no longer \ + matches the schema. Rebuild the segment against the current schema.", + dataset.manifest.version + ))); + } + } let stale_fragments = segment .fragment_bitmap() .iter() @@ -285,9 +310,9 @@ async fn prune_stale_segment_coverage( return true; }; let historical_files = - fragment_field_files(&historical, historical_fragment, &indexed_field_ids); + fragment_field_files(&historical, historical_fragment, &relevant_field_ids); let current_files = - fragment_field_files(dataset, current_fragment, &indexed_field_ids); + fragment_field_files(dataset, current_fragment, &relevant_field_ids); let changed_files = historical_files.is_none() || historical_files != current_files; let changed_overlays = prune_newer_overlays @@ -297,7 +322,7 @@ async fn prune_stale_segment_coverage( .data_file .fields .iter() - .any(|field_id| indexed_field_ids.contains(field_id)) + .any(|field_id| relevant_field_ids.contains(field_id)) }); changed_files || changed_overlays }) @@ -3596,6 +3621,18 @@ impl DatasetIndexInternalExt for Dataset { ))); } + // `field_paths` is already the keyed prefix -- it was built from `keyed_fields()` + // above -- so a carried column is deliberately NOT resolved or type-checked here. + // That is the point: the rebuild writes fresh, non-covered metadata and never reads + // the carried payload, so demanding those columns exist in the target would reject + // a target that can hold this index perfectly well. + // + // `initialize_vector_index` then re-resolves the carried columns itself, by name, + // against the target schema. `initialize_scalar_index` does NOT -- it recreates + // with default params and never reads `covering_fields`, so a covered scalar index + // would land on the target with covering silently dropped. Unreachable today: + // `covering_columns` exists only on `VectorIndexParams`, so nothing produces a + // covered scalar index. Whoever adds a scalar producer must fix that path too. let field_names = field_paths.iter().map(String::as_str).collect::>(); if let Some(index_details) = &source_index.index_details { let index_details_wrapper = IndexDetails(index_details.clone()); @@ -6526,6 +6563,90 @@ mod tests { assert_eq!(desc.rows_indexed(), 10); } + /// `describe_indices` with a column criteria must still find a covered index. + /// + /// `index_matches_criteria` is documented to receive only the column named by + /// `for_column` -- it rejects anything else with `fields.len() != 1`. A covered + /// index lists its carried columns in `fields` too, so resolving all of `fields` + /// silently fails that check and drops the index from the description with no + /// error: `describe_indices(for_column)` reports "no index on this column" for a + /// column that is indexed. + #[tokio::test] + async fn test_describe_indices_for_column_finds_covered_index() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(20), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("covered_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + let covered = dataset.load_indices_by_name("covered_idx").await.unwrap(); + assert_eq!(covered[0].fields.len(), 2, "index should be covered"); + + let descriptions = dataset + .describe_indices(Some(IndexCriteria::default().for_column("vector"))) + .await + .unwrap(); + assert_eq!( + descriptions.len(), + 1, + "a covered index must still be described for its keyed column" + ); + assert_eq!(descriptions[0].name(), "covered_idx"); + + // The carried column is not the indexed column, so it must not match. + let by_carried = dataset + .describe_indices(Some(IndexCriteria::default().for_column("id"))) + .await + .unwrap(); + assert!( + by_carried.is_empty(), + "a carried column is not indexed and must not match `for_column`" + ); + + // `field_ids()` is a capability statement -- "the fields the index is built on" -- + // so it must report the keyed column only. Reporting the carried column here tells + // every consumer (Python `describe_indices().fields`, Java `getFieldIds()`, the + // namespace REST `columns` list) that this index can serve `id`; Python's + // `_default_vector_index_for_column` then hands back `vec`'s IVF centroids for the + // int32 column `id` instead of raising. + let vector_field_id = dataset.schema().field("vector").unwrap().id as u32; + let id_field_id = dataset.schema().field("id").unwrap().id as u32; + let desc = &descriptions[0]; + assert_eq!( + desc.field_ids(), + &[vector_field_id], + "field_ids() must report the keyed column only, not the carried one" + ); + // The carried set is not lost -- it stays reachable per segment. + assert_eq!( + desc.metadata()[0].covering_fields, + vec![id_field_id as i32], + "the carried column must remain reachable via metadata()[i].covering_fields" + ); + } + #[tokio::test] async fn test_describe_indices_derives_type_from_url_without_plugin() { // When index details exist but no plugin is registered for the type @@ -7918,6 +8039,68 @@ mod tests { ); } + /// `initialize_index` must copy a covered index, not reject it. + /// + /// It derives the column names it initializes from `source_index.fields`, and a + /// covered index lists its carried columns there too. Both initializers accept + /// exactly one name (they re-resolve the carried columns themselves from + /// `covering_fields`), so passing all of `fields` makes every covered index + /// un-copyable -- and `initialize_indices` then fails the whole dataset copy. + #[tokio::test] + async fn test_initialize_index_copies_covered_index() { + use arrow_array::types::Float32Type; + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/{}", test_dir, "source"); + let target_uri = format!("{}/{}", test_dir, "target"); + + let make_reader = || { + lance_datagen::gen_batch() + .col("id", array::step::()) + .col("vector", array::rand_vec::(8.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)) + }; + + let mut source_dataset = Dataset::write(make_reader(), &source_uri, None) + .await + .unwrap(); + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 4, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["id".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("covered_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + + let mut target_dataset = Dataset::write(make_reader(), &target_uri, None) + .await + .unwrap(); + target_dataset + .initialize_index(&source_dataset, "covered_idx") + .await + .unwrap(); + + let target_indices = target_dataset.load_indices().await.unwrap(); + assert_eq!(target_indices.len(), 1); + let id_field_id = target_dataset.schema().field("id").unwrap().id; + assert_eq!( + target_indices[0].covering_fields, + vec![id_field_id], + "the copied index must keep its covering declaration" + ); + target_indices[0].validate_covering_fields().unwrap(); + } + #[rstest] #[case::simple("value", "data.value")] #[case::quoted("value.with.dot", "data.`value.with.dot`")] @@ -8790,6 +8973,377 @@ mod tests { ); } + /// A covering segment's `covering_fields` must survive `commit_existing_index_segments`. + /// The segment files still hold the payload, so dropping the declaration would silently + /// fall covered queries back to a base-table take. + #[tokio::test] + async fn test_commit_existing_index_segments_preserves_covering_fields() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(20), BatchCount::from(2)); + + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 10, + max_rows_per_group: 10, + ..Default::default() + }), + ) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + // A real covered segment: its auxiliary storage physically holds the "id" + // payload, which the commit-boundary cross-check verifies. + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["id".to_string()]); + let seg = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg.covering_fields, vec![id_field_id]); + + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg)], + ) + .await + .unwrap(); + + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + assert_eq!( + committed[0].covering_fields, + vec![id_field_id], + "committed segment must preserve its covering `covering_fields`" + ); + } + + /// A second commit whose incoming segment covers a different column set than an + /// existing, disjoint segment we would *retain* must be rejected: all segments of one + /// logical index must agree on their covering, or the committed `covering_fields` + /// would misdescribe some segment's storage. + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_inconsistent_retained_covering() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col( + "vector", + array::rand_vec::(8.into()), + ) + .into_reader_rows(RowCount::from(20), BatchCount::from(1)); + + let mut dataset = Dataset::write( + reader, + test_uri, + Some(WriteParams { + max_rows_per_file: 10, + max_rows_per_group: 10, + ..Default::default() + }), + ) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + + // First commit: a real segment covering `id`, over fragment 0 only. + let mut covered_params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + covered_params.covering_columns(vec!["id".to_string()]); + let seg0 = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .fragments(vec![0]) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg0.covering_fields, vec![id_field_id]); + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg0)], + ) + .await + .unwrap(); + + // Second commit: a disjoint segment (fragment 1) that covers *nothing*. It would + // be retained alongside seg0, whose covering is `[id]` -- an inconsistency. + let plain_params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + let seg1 = dataset + .create_index_builder(&["vector"], IndexType::Vector, &plain_params) + .name("vector_idx".to_string()) + // The name already exists (seg0 was committed); replace(true) only skips + // the duplicate-name check for this uncommitted build. + .replace(true) + .fragments(vec![1]) + .execute_uncommitted() + .await + .unwrap(); + assert!(seg1.covering_fields.is_empty()); + let result = dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg1)], + ) + .await; + let err = + result.expect_err("mismatched covering across retained segments must be rejected"); + assert!( + err.to_string() + .contains("a logical index cannot mix declarations") + && err.to_string().contains("covering_fields"), + "expected a covering-consistency rejection, got: {err}" + ); + } + + /// When a segment covers a column whose data file was rewritten in place after the + /// segment was built, `prune_stale_segment_coverage` must drop the affected fragment + /// from the segment's coverage -- even though the *indexed* column's file is untouched. + /// Otherwise a covered query reads the segment's stale copy of that column. + #[tokio::test] + async fn test_commit_existing_index_segments_prunes_stale_covered_fragment() { + use crate::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let dim = 8i32; + let n = 20i32; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("tag", DataType::Int32, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n * dim).map(|v| v as f32).collect::>()), + dim, + ) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n).collect::>())), + Arc::new(Int32Array::from((0..n).collect::>())), + Arc::new(vectors), + ], + ) + .unwrap(); + // Two fragments (ids 0..10 -> frag 0, 10..20 -> frag 1). + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + test_uri, + Some(WriteParams { + max_rows_per_file: 10, + max_rows_per_group: 10, + ..Default::default() + }), + ) + .await + .unwrap(); + + // A scalar index on the merge key routes the partial-schema update through the + // in-place column-rewrite path (RewriteColumns) instead of a row-move, so the + // covered column's file changes while the fragment id is preserved. + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let tag_field_id = dataset.schema().field("tag").unwrap().id; + + // Build a real covering segment BEFORE the update, so it is stamped at the + // pre-update version and covers `tag`, whose files still hold the old values. + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["tag".to_string()]); + let seg = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg.covering_fields, vec![tag_field_id]); + + // In-place rewrite of the covered `tag` column for fragment 0 only (no covering + // index committed yet, so this patches the column file in place and keeps the + // fragment id). The indexed `vector` column's file is untouched. + let sub_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("tag", DataType::Int32, true), + ])); + let update = RecordBatch::try_new( + sub_schema.clone(), + vec![ + Arc::new(Int32Array::from((0..5).collect::>())), + Arc::new(Int32Array::from((100..105).collect::>())), + ], + ) + .unwrap(); + let (updated, _) = MergeInsertBuilder::try_new(Arc::new(dataset), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap() + .execute_reader(Box::new(RecordBatchIterator::new([Ok(update)], sub_schema))) + .await + .unwrap(); + let mut dataset = Arc::try_unwrap(updated).unwrap_or_else(|arc| (*arc).clone()); + + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg)], + ) + .await + .unwrap(); + + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed.len(), 1); + let bitmap = committed[0].fragment_bitmap.as_ref().unwrap(); + assert!( + !bitmap.contains(0), + "fragment 0's covered column was rewritten; its stale coverage must be pruned" + ); + assert!( + bitmap.contains(1), + "fragment 1's covered column was untouched; it stays covered" + ); + } + + /// A covered *struct* that grows a child (via `add_columns`) after a segment was built + /// makes the segment's stored covering payload type-incompatible for every covered row -- + /// a global change no per-fragment prune can repair. Since the covering index is not + /// committed while the segment is staged, the Merge-commit guard never sees the add; + /// `prune_stale_segment_coverage` must reject the segment commit instead. + #[tokio::test] + async fn test_commit_existing_index_segments_rejects_grown_covered_struct() { + use arrow_array::StructArray; + use arrow_schema::Fields; + use lance_file::version::LanceFileVersion; + + use crate::dataset::NewColumnTransform; + + let test_dir = tempfile::tempdir().unwrap(); + let test_uri = test_dir.path().to_str().unwrap(); + + let dim = 8i32; + let n = 10i32; + let meta_fields = Fields::from(vec![Field::new("a", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("meta", DataType::Struct(meta_fields.clone()), true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let a = Arc::new(Int32Array::from((0..n).collect::>())); + let meta = Arc::new(StructArray::new(meta_fields, vec![a], None)); + let vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from((0..n * dim).map(|v| v as f32).collect::>()), + dim, + ) + .unwrap(); + let batch = RecordBatch::try_new(schema.clone(), vec![meta, Arc::new(vectors)]).unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let meta_field_id = dataset.schema().field("meta").unwrap().id; + + // Build a real covering segment stamped at the current (pre-add) version, + // covering the struct column `meta`. + let mut params = crate::index::vector::VectorIndexParams::ivf_flat( + 1, + lance_linalg::distance::MetricType::L2, + ); + params.covering_columns(vec!["meta".to_string()]); + let seg = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(seg.covering_fields, vec![meta_field_id]); + + // Grow the covered struct AFTER the segment was built. No covering index is + // committed yet, so the Merge-commit guard does not fire and the add succeeds. + let add = + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![Field::new( + "meta", + DataType::Struct(Fields::from(vec![Field::new("b", DataType::Int32, true)])), + true, + )]))); + dataset.add_columns(add, None, None).await.unwrap(); + + // Committing the now-stale segment must be rejected: its stored covering payload + // still holds the pre-child struct type. + let err = dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&seg)], + ) + .await + .expect_err("committing a segment covering a since-grown struct must be rejected"); + assert!( + err.to_string().contains("changed its subtree"), + "expected a covered-subtree-change rejection, got: {err}" + ); + } + #[tokio::test] async fn test_commit_existing_index_segments_rejects_duplicate_segment_ids() { use lance_datagen::{BatchCount, RowCount, array}; diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 92c3655827f..35654acafbe 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -4002,6 +4002,332 @@ mod tests { assert_eq!(got, vec![10, 11, 12, 13]); } + /// The covered recovery takes its payload with `MissingRowPolicy::Ignore`, which silently + /// returns fewer rows than requested when an id no longer resolves. A stale prefilter can + /// list such ids: a scalar index built over a fragment that a later delete removed keeps + /// emitting that fragment's rows, and `create_deletion_mask_impl` produces no mask at all + /// when every fragment in the VECTOR index's bitmap is intact. Pairing the requested ids + /// with the returned payload then fails with "all columns in a record batch must have the + /// same length" -- a query the same dataset answers fine without covering. + /// + /// Scope: this covers the *unresolvable-id* shortfall only -- the fragment is gone, so + /// `get_row_addrs` drops the ids before the read. The other way a take can come up short -- + /// an id that resolves to a live address whose row is tombstoned -- is not exercised here, + /// and is not reachable. + /// + /// Note those are two different masks, despite the shared name. The one above is + /// `create_deletion_mask_impl`'s, scoped to the vector index's bitmap, and it can indeed be + /// absent. The one that makes tombstoned rows unreachable is `DatasetPreFilter`'s + /// `deleted_ids`, which is intersected in separately and spans every fragment in the manifest, + /// so + /// it holds no matter which `PreFilterSource` produced the ids. + #[tokio::test] + async fn test_covered_recovery_tolerates_unresolvable_prefilter_ids() { + use arrow_array::{Int32Array, RecordBatchIterator}; + + let dim = 4i32; + let frag0 = 16usize; + let extra = 3usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + // Two separated clusters so `early_pruning` leaves an unsearched partition and the + // late-search shortcut (which arms the covered recovery) is reachable. + let make_batch = |ids: Vec| { + let values: Vec = ids + .iter() + .flat_map(|id| { + let center = if id % 2 == 0 { 0.0f32 } else { 1000.0 }; + (0..dim as usize).map(move |d| center + (*id as usize * 4 + d) as f32 * 1e-3) + }) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim) + .unwrap(), + ), + ], + ) + .unwrap() + }; + + let mut dataset = Dataset::write( + RecordBatchIterator::new( + [Ok(make_batch((0..frag0 as i32).collect()))], + schema.clone(), + ), + "memory://covered_recovery_stale_prefilter", + Some(WriteParams { + enable_stable_row_ids: true, + max_rows_per_file: frag0, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Covered index over fragment 0 only. + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + // Fragment 1, then a scalar index spanning BOTH fragments. + let ids: Vec = (frag0 as i32..(frag0 + extra) as i32).collect(); + dataset + .append( + RecordBatchIterator::new([Ok(make_batch(ids))], schema.clone()), + None, + ) + .await + .unwrap(); + let scalar = lance_index::scalar::ScalarIndexParams::for_builtin( + lance_index::scalar::BuiltinIndexType::BTree, + ); + dataset + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + // Drop fragment 1 entirely. The BTree still lists its ids, so the prefilter admits + // rows the row-id index can no longer resolve. + dataset.delete("id >= 16").await.unwrap(); + + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 10).unwrap(); + scan.minimum_nprobes(1); + scan.filter("id >= 15").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + + let batch = scan + .try_into_batch() + .await + .expect("covered query must tolerate prefilter ids that no longer resolve"); + let ids = batch["id"].as_primitive::(); + let got: Vec = ids.values().to_vec(); + assert_eq!( + got, + vec![15], + "only the surviving admitted row may come back; ids from the deleted fragment \ + must be dropped, not paired with mismatched payload" + ); + } + + /// The covered end-of-stream recovery re-emits every prefilter row the search did not + /// emit. `DatasetPreFilter` produces a bounded ALLOW LIST from deletions alone on a + /// stable-row-id dataset, so that recovery fires on queries carrying NO filter at all + /// (`PreFilterSource::None`). Tracking the emitted row ids must therefore not be gated + /// on a prefilter source being present -- otherwise the "already emitted" set is empty + /// and every live row is emitted a second time with INFINITY distance. + #[tokio::test] + async fn test_covered_unfiltered_query_does_not_duplicate_rows() { + use arrow_array::types::{Int32Type, UInt64Type}; + use arrow_array::{Int32Array, RecordBatchIterator}; + + const INDEX: &str = "vec_idx"; + let dim = 4i32; + let n = 12usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + // Two well-separated clusters with the query sitting on top of the first: the second + // centroid is orders of magnitude farther, so `early_pruning` leaves minimum_nprobes + // at 1. An unsearched partition is what makes the late-search no-rows shortcut (and + // hence the covered recovery) reachable -- with evenly spread vectors early pruning + // raises minimum_nprobes to cover every partition and the shortcut never fires. + let ids: Vec = (0..n as i32).collect(); + let values: Vec = (0..n) + .flat_map(|r| { + let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; + (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) + }) + .collect(); + let vector = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Float32, true)), + dim, + Arc::new(Float32Array::from(values)), + None, + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(vector)], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covered_unfiltered_no_dupes", + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // A deletion makes the stable-row-id deletion mask a bounded ALLOW LIST, so + // `max_len()` is Some(live) even though the query below carries no filter. + dataset.delete("id = 1").await.unwrap(); + let live = n - 1; + + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + // k > live so the "fewer than k prefilter matches" shortcut condition holds. + scan.nearest("vector", &q, live + 9).unwrap(); + scan.minimum_nprobes(1); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + let batch = scan.try_into_batch().await.unwrap(); + + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let mut seen: Vec = row_ids.values().to_vec(); + seen.sort_unstable(); + let mut distinct = seen.clone(); + distinct.dedup(); + assert_eq!( + seen.len(), + distinct.len(), + "covered unfiltered query returned DUPLICATE row ids: {seen:?}" + ); + assert_eq!( + batch.num_rows(), + live, + "covered unfiltered query must return each live row exactly once" + ); + + let ids = batch + .column_by_name("id") + .expect("covered 'id' must be emitted") + .as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i) as u64, + row_ids.value(i), + "covered id must stay row-aligned" + ); + } + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + (0..n as i32).filter(|id| *id != 1).collect::>(), + "every live row must appear exactly once" + ); + } + + /// A covered ("included") struct's payload schema is fixed at index build time. Growing + /// that struct via `add_columns` (which commits as `Operation::Merge`) -- even an + /// AllNulls, metadata-only child that writes no data file -- would leave the index + /// emitting the old struct type while covered queries declare the new one, an Arrow type + /// mismatch. The commit boundary must reject it (drop the index first), mirroring the + /// `Project` drop/alter guard. + #[tokio::test] + async fn test_add_columns_child_to_covered_struct_is_rejected() { + use arrow_array::{Int32Array, RecordBatchIterator, StructArray}; + use arrow_schema::Fields; + use lance_file::version::LanceFileVersion; + + use crate::dataset::NewColumnTransform; + + let n = 64usize; + let dim = 4i32; + let meta_fields = Fields::from(vec![Field::new("a", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("meta", DataType::Struct(meta_fields.clone()), true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + + let a = Arc::new(Int32Array::from((0..n as i32).collect::>())); + let meta = Arc::new(StructArray::new(meta_fields, vec![a], None)); + let values: Vec = (0..n * dim as usize).map(|i| i as f32 + 1.0).collect(); + let vector = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim).unwrap(), + ); + let batch = RecordBatch::try_new(schema.clone(), vec![meta, vector]).unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema.clone()), + "memory://covered_struct_add_child", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["meta".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vec_idx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Add child `meta.b` -- merges into the covered struct, changing its type. + let add = + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![Field::new( + "meta", + DataType::Struct(Fields::from(vec![Field::new("b", DataType::Int32, true)])), + true, + )]))); + let err = dataset + .add_columns(add, None, None) + .await + .expect_err("growing a covered struct's subtree must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("included") && msg.contains("vec_idx"), + "expected a covered-field rejection naming the index, got: {err}" + ); + } + /// Read-side payoff for every vector index type: a query projecting only a covered /// column is satisfied from the index -- no `TakeExec` against the base table -- /// with row-aligned values and sane recall. @@ -4824,21 +5150,39 @@ mod tests { assert_eq!(batch.num_rows(), 5); assert!(batch.column_by_name("id").is_some()); - // The hard case: a bounded selective prefilter mixes the plan's 2-column - // not-found shortcut batch with search batches -- if the search batches still - // carried the undeclared payload column, the widths would disagree where the - // stream is concatenated and the query would fail. - let mut scan = dataset.scan(); - scan.nearest("vector", &q, 5).unwrap(); - scan.filter("id < 3").unwrap(); - scan.prefilter(true); - scan.project(&["id"]).unwrap(); - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!( - batch.num_rows(), - 3, - "all prefilter-matched rows must come back" - ); + // The hard case: a bounded selective prefilter over unsearched partitions makes + // `late_search` emit its 2-column not-found shortcut batch into the same stream as + // the search batches. Unnarrowed, those search batches still carry the undeclared + // payload column, so the widths disagree where DataFusion's TopK concatenates them + // -- `interleave_record_batch` reads every batch through the FIRST batch's schema + // and panics with an out-of-bounds column index. + // Repeated because the unnarrowed failure is order-dependent: it only panics when + // TopK happens to store a wide batch first, which is roughly 4 runs in 5. One + // iteration would make this test pass intermittently against a real regression. + for attempt in 0..5 { + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.minimum_nprobes(1); + scan.maximum_nprobes(NUM_CLUSTERS); + scan.filter("id < 3").unwrap(); + scan.prefilter(true); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + 3, + "all prefilter-matched rows must come back (attempt {attempt})" + ); + let ids = batch["id"].as_primitive::(); + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + vec![0, 1, 2], + "covered payload must be taken from the base table once the declaration is \ + gone (attempt {attempt})" + ); + } } /// A covered index that searches zero partitions (empty heap) diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 4a824058640..36b2b03663b 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -10,7 +10,7 @@ use crate::{ dataset::transaction::{DataOverlayGroup, Operation, Transaction, UpdateMode}, }; use futures::{StreamExt, TryStreamExt}; -use lance_core::{Error, Result, utils::deletion::DeletionVector}; +use lance_core::{Error, Result, datatypes::Schema, utils::deletion::DeletionVector}; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{CompactedSsTable, MEM_WAL_INDEX_NAME}; use lance_select::{RowAddrTreeMap, RowSetOps}; @@ -37,6 +37,57 @@ pub struct TransactionRebase<'a> { /// Compacted SSTables from conflicting UpdateMemWalState transactions. /// Used when rebasing CreateIndex of MemWalIndex. conflicting_mem_wal_compacted_sstables: Vec, + /// The dataset schema at the read version, used to expand a covering index's + /// covered (included) struct fields to their leaf subtrees when checking whether a + /// concurrent transaction touches covered data. + schema: Arc, + /// Read-version snapshot of which data file backs each covered (included) leaf + /// field, per fragment: `fragment id -> field id -> path`. Empty unless this + /// transaction creates an index that declares covering columns. + /// + /// This is what lets `check_create_index_txn` tell a concurrent `Merge` that + /// rewrote covered data *in place* from one that merely added a column. The + /// post-merge fragments alone cannot answer that: a rewrite tombstones the + /// superseded field with a single sentinel that does not name it. + covered_field_paths: Arc>>, +} + +/// Snapshot the data file backing each covered leaf field of `new_indices`, as of +/// `dataset`'s version. Covered struct fields are expanded to their whole subtree +/// because a rewritten data file lists leaf ids while `covering_fields` records only +/// the parent id. Fragments with no covered field are omitted. +fn covered_field_paths_at( + dataset: &Dataset, + new_indices: &[IndexMetadata], + schema: &Schema, +) -> HashMap> { + let mut covered: HashSet = HashSet::new(); + for index in new_indices { + for &id in index.covering_fields.iter() { + match schema.field_by_id(id) { + Some(field) => crate::index::collect_subtree_field_ids(field, &mut covered), + None => { + covered.insert(id); + } + } + } + } + if covered.is_empty() { + return HashMap::new(); + } + dataset + .manifest + .fragments + .iter() + .filter_map(|fragment| { + let paths: HashMap = Transaction::fragment_field_paths(fragment) + .into_iter() + .filter(|(field_id, _)| covered.contains(field_id)) + .map(|(field_id, path)| (field_id, path.to_string())) + .collect(); + (!paths.is_empty()).then_some((fragment.id, paths)) + }) + .collect() } /// Whether `operation` may make a nullability-affecting schema change: a @@ -99,6 +150,55 @@ impl<'a> TransactionRebase<'a> { transaction: Transaction, affected_rows: Option<&'a RowAddrTreeMap>, ) -> Result { + // Capture the schema AS OF THE TRANSACTION'S READ VERSION, matching what + // `initial_fragments_for_rebase` does for fragments. The dataset handed to us is not + // necessarily at that version: the URI commit path deliberately loads the *latest* + // version ("we are writing to the main history, and need to check out the latest + // version") for any non-detached read version. The covering checks below use this as + // the *pre-commit* side when asking whether a concurrent operation changed a covered + // field's subtree, so capturing the latest schema would compare a post-change schema + // against itself -- silently missing a covered-field rename or retype and committing + // index storage under a schema it was never built against. + // + // Only `check_create_index_txn` and `check_data_replacement_txn` consult this, so + // only they pay for the extra checkout; every other operation gets the in-hand schema + // and never reads it. If a new check starts using `self.schema`, add its operation + // here or it will silently see the wrong version. + // + // CreateIndex is narrowed further to indices that actually declare covering columns: + // comparing a covered subtree across versions is the only thing that can observe the + // difference, so an ordinary index build keeps the pre-covering behaviour (in-hand + // schema, no extra round trip) instead of paying for a manifest load on every commit + // retry. DataReplacement cannot be narrowed the same way -- the covering it has to + // respect belongs to a *concurrent* transaction's indices, which are not known until + // the check runs. + // + // `read_version == 0` is the "no prior version" sentinel (overwrite/create): there is + // nothing to check out, and attempting it fails with DatasetNotFound. + let needs_read_version_schema = match &transaction.operation { + Operation::CreateIndex { new_indices, .. } => new_indices + .iter() + .any(|idx| !idx.covering_fields.is_empty()), + Operation::DataReplacement { .. } => true, + _ => false, + }; + let checked_out; + let at_read_version = if !needs_read_version_schema + || transaction.read_version == 0 + || dataset.manifest.version == transaction.read_version + { + dataset + } else { + checked_out = dataset.checkout_version(transaction.read_version).await?; + &checked_out + }; + let schema = Arc::new(at_read_version.schema().clone()); + let covered_field_paths = Arc::new(match &transaction.operation { + Operation::CreateIndex { new_indices, .. } if needs_read_version_schema => { + covered_field_paths_at(at_read_version, new_indices, &schema) + } + _ => HashMap::new(), + }); match &transaction.operation { // These operations add new fragments or don't modify any. Operation::Append { .. } @@ -117,6 +217,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids: HashSet::new(), conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }), Operation::Delete { updated_fragments, @@ -145,6 +247,8 @@ impl<'a> TransactionRebase<'a> { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }); } @@ -158,6 +262,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::Rewrite { groups, .. } => { @@ -176,6 +282,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::DataReplacement { replacements } => { @@ -191,6 +299,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::DataOverlay { groups } => { @@ -206,6 +316,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } Operation::Merge { fragments, .. } => { @@ -220,6 +332,8 @@ impl<'a> TransactionRebase<'a> { modified_fragment_ids, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + schema: schema.clone(), + covered_field_paths: covered_field_paths.clone(), }) } } @@ -695,6 +809,10 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { + // Captured before the mutable borrow below so the covered-field checks can expand + // covered struct fields to their leaf subtree (an Arc clone is cheap). + let schema = self.schema.clone(); + let covered_field_paths = self.covered_field_paths.clone(); if let Operation::CreateIndex { new_indices, removed_indices, @@ -752,25 +870,89 @@ impl<'a> TransactionRebase<'a> { fields_modified, .. } => { + // Expand covered struct fields to their leaf subtree so a concurrent + // update of a covered struct's child still invalidates the fragment's + // coverage (the schema is captured at the read version above). Transaction::prune_updated_fields_from_indices( new_indices, updated_fragments, fields_modified, + schema.as_ref(), ); Ok(()) } - // Merge, reserve, and project don't change row ids. The MemWAL - // index is the exception: its install validates schema-dependent - // state, which a concurrent schema change invalidates. - Operation::Merge { .. } => { - if new_indices.iter().any(|idx| idx.name == MEM_WAL_INDEX_NAME) { + // Reserve and project don't change row ids. Merge is usually fine too, + // but two independent conditions make it unsafe here, and neither + // subsumes the other: + // + // (a) the MemWAL index's install validates schema-dependent state, which + // a concurrent schema change invalidates. + // + // (b) it can rewrite a covered column's data file *in place*, or change a + // covered field's *subtree* (notably `add_columns` growing a covered + // struct with an AllNulls child, which writes NO data file, so the + // path comparison alone is blind to it). Either leaves the rebased + // index serving values it was not built against. + // + // Merely adding an unrelated column trips none of these and keeps the + // optimistic path -- otherwise a routine `add_columns` would discard a + // covered index build that may have been running for hours. + Operation::Merge { + fragments, + schema: merged_schema, + .. + } => { + let installs_mem_wal_index = + new_indices.iter().any(|idx| idx.name == MEM_WAL_INDEX_NAME); + let rewrote_covered_data = fragments.iter().any(|fragment| { + let Some(previous) = covered_field_paths.get(&fragment.id) else { + return false; + }; + let current = Transaction::fragment_field_paths(fragment); + previous.iter().any(|(field_id, previous_path)| { + current + .get(field_id) + .is_some_and(|current_path| *current_path != previous_path.as_str()) + }) + }); + let changed_covered_subtree = new_indices.iter().any(|idx| { + idx.covering_fields.iter().any(|id| { + Transaction::covered_field_subtree_changed(&schema, merged_schema, *id) + }) + }); + if installs_mem_wal_index || rewrote_covered_data || changed_covered_subtree { Err(self.retryable_conflict_err(other_transaction, other_version)) } else { Ok(()) } } Operation::ReserveFragments { .. } => Ok(()), - Operation::Project { .. } => Ok(()), + Operation::Project { + schema: projected_schema, + .. + } => { + // A concurrent Project that changes a covered field's subtree in *any* + // way -- drop, rename, retype, nullability, child change -- leaves the + // rebased index emitting a schema the dataset no longer declares, so + // force a retry and let the index rebuild against the projected schema. + // A rename counts: `covered_field_subtree_changed` compares names, and + // the index's covering payload carries the old one. + // Non-covering indexes, and covered subtrees the Project left alone, + // keep the optimistic path (row ids are unchanged either way). + if new_indices.iter().any(|idx| { + idx.covering_fields.iter().any(|id| { + Transaction::covered_field_subtree_changed( + &schema, + projected_schema, + *id, + ) + }) + }) { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } // Should be compatible with rewrite if it didn't move the rows // we indexed. If it did, we could retry. // TODO: this will change with stable row ids. @@ -848,16 +1030,19 @@ impl<'a> TransactionRebase<'a> { } Operation::UpdateConfig { .. } => Ok(()), Operation::DataReplacement { replacements } => { - // A data replacement only conflicts if it is updating a field the - // index depends on -- whether keyed on or merely carried, since - // `fields` lists both (see `IndexMetadata::covering_fields`). - let newly_depended_fields = new_indices - .iter() - .flat_map(|idx| idx.fields.iter()) - .collect::>(); + // A data replacement conflicts if it updates a field the index + // either *indexes* or *covers*: a covered field's values are + // materialized in the index storage, so replacing them would + // leave the rebased index serving stale values. Expand each field to + // its leaf subtree, since `covering_fields` records a covered struct's + // parent id while a replacement lists leaf ids. + let mut relevant_fields: HashSet = HashSet::new(); + for idx in new_indices.iter() { + relevant_fields.extend(Transaction::index_dependent_leaf_ids(idx, &schema)); + } for replacement in replacements { for field in replacement.1.fields.iter() { - if newly_depended_fields.contains(&field) { + if relevant_fields.contains(field) { return Err( self.retryable_conflict_err(other_transaction, other_version) ); @@ -1169,6 +1354,7 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { + let schema = self.schema.clone(); if let Operation::DataReplacement { replacements } = &self.transaction.operation { match &other_transaction.operation { Operation::Append { .. } @@ -1268,13 +1454,17 @@ impl<'a> TransactionRebase<'a> { // the index's fragment bitmap, which would lead to fewer conflicts. However // this would introduce fragment bitmaps with holes which may not be well tested // yet. For now, we don't allow this case. - let newly_depended_fields = new_indices - .iter() - .flat_map(|idx| idx.fields.iter()) - .collect::>(); + // Expand covered fields to their leaf subtree so replacing a covered + // struct's subfield (a leaf id, while `covering_fields` records the + // parent) is still recognized as touching the index -- matching the + // reverse rebase direction in check_create_index_txn. + let mut relevant_fields: HashSet = HashSet::new(); + for idx in new_indices.iter() { + relevant_fields.extend(Transaction::index_dependent_leaf_ids(idx, &schema)); + } for replacement in replacements { for field in replacement.1.fields.iter() { - if newly_depended_fields.contains(&field) { + if relevant_fields.contains(field) { return Err( self.retryable_conflict_err(other_transaction, other_version) ); @@ -2329,6 +2519,14 @@ mod tests { }; use lance_table::format::DataFile; + /// An empty schema for manually-constructed rebases in tests: covered-field subtree + /// expansion is a no-op against it (each id maps to itself), preserving the exact-id + /// behavior these tests assert. Tests that exercise covered structs build the rebase + /// via `try_new`, which captures the real dataset schema. + fn empty_lance_schema() -> Arc { + Arc::new(lance_core::datatypes::Schema::default()) + } + async fn test_dataset(num_rows: usize, num_fragments: usize) -> Dataset { let write_params = WriteParams { max_rows_per_file: num_rows / num_fragments, @@ -2690,6 +2888,279 @@ mod tests { assert!(matches!(err, Error::DatasetNotFound { .. })); } + /// A concurrent Project that drops a column an in-flight index *covers* must force + /// the index build to rebase (retryable conflict); a Project that preserves the + /// covered column stays on the optimistic path. + #[tokio::test] + async fn test_create_index_conflicts_with_project_dropping_covered_field() { + use roaring::RoaringBitmap; + + let dataset = test_dataset(10, 2).await; + let a_id = dataset.schema().field("a").unwrap().id; + let b_id = dataset.schema().field("b").unwrap().id; + + // A covering index over key "a" that also covers "b". + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![a_id, b_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32, 1])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![b_id], + }; + let self_txn = Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![covering_index], + removed_indices: vec![], + }, + ); + + // Concurrent Project that drops "b" (the covered column) -> conflict. + let dropped_schema = dataset.schema().project(&["a"]).unwrap(); + let drop_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: dropped_schema, + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn.clone(), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&drop_txn, 2).is_err(), + "dropping a covered column must conflict with an in-flight covering index build" + ); + + // A Project that preserves "b" -> no conflict. + let keep_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: dataset.schema().clone(), + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn.clone(), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&keep_txn, 2).is_ok(), + "a project that preserves the covered column must not conflict" + ); + + // A Project that renames the covered column (id stays, name changes) -> conflict: + // the built index still emits the old name. + let mut renamed_schema = dataset.schema().clone(); + renamed_schema.field_by_id_mut(b_id).unwrap().name = "b_renamed".to_string(); + let rename_txn = Transaction::new_from_version( + 2, + Operation::Project { + schema: renamed_schema, + preserves_nullability: true, + }, + ); + let mut rebase = TransactionRebase::try_new(&dataset, self_txn, None) + .await + .unwrap(); + assert!( + rebase.check_txn(&rename_txn, 2).is_err(), + "renaming a covered column must conflict with an in-flight covering index build" + ); + } + + /// A concurrent DataReplacement of a *child* of a covered struct must conflict with + /// an in-flight covering index build: `covering_fields` records the struct's parent + /// id, but the replacement lists the leaf id, so the check must expand the subtree. + #[tokio::test] + async fn test_create_index_conflicts_with_replacement_of_covered_struct_child() { + use arrow_schema::{ + DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, + }; + use roaring::RoaringBitmap; + + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])), + true, + ), + ]); + let schema = lance_core::datatypes::Schema::try_from(&arrow).unwrap(); + let vec_id = schema.field("vec").unwrap().id; + let s_id = schema.field("s").unwrap().id; + let a_id = schema.field("s.a").unwrap().id; + + // Index keyed on `vec`, covering the whole struct `s`. + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![vec_id, s_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![s_id], + }; + let self_txn = Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![covering_index], + removed_indices: vec![], + }, + ); + + // try_new needs a real dataset, so build the rebase directly with the struct schema. + let mut rebase = TransactionRebase { + transaction: self_txn, + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::new(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(schema), + }; + + // A concurrent DataReplacement rewriting the covered struct's child `s.a`. + let replace_child = Transaction::new_from_version( + 2, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("child.lance", vec![a_id], None), + )], + }, + ); + assert!( + rebase.check_txn(&replace_child, 2).is_err(), + "replacing a child of a covered struct must conflict (subtree expansion)" + ); + } + + /// A concurrent `add_columns` (which commits as `Operation::Merge`) that grows a + /// covered struct with an AllNulls child writes **no data file**, so the covered-file + /// path comparison cannot see it. The index would otherwise commit storage built for + /// the old struct type under the new schema, and every covered query on it would fail. + /// The Merge arm must therefore also compare the covered subtree, as the Project arm + /// does -- the two checks are complementary and neither subsumes the other. + #[tokio::test] + async fn test_create_index_conflicts_with_merge_growing_covered_struct() { + use arrow_schema::{ + DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, + }; + use roaring::RoaringBitmap; + + let struct_of = |children: Vec| { + ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("s", DataType::Struct(ArrowFields::from(children)), true), + ]) + }; + let before = lance_core::datatypes::Schema::try_from(&struct_of(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])) + .unwrap(); + // The AllNulls child add keeps every existing id and gives the new leaf a fresh one. + let after = lance_core::datatypes::Schema::try_from(&struct_of(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])) + .unwrap(); + let vec_id = before.field("vec").unwrap().id; + let s_id = before.field("s").unwrap().id; + assert_eq!( + s_id, + after.field("s").unwrap().id, + "the struct keeps its id" + ); + + // `covering_fields` is always the trailing entries of `fields`. + let covering_index = IndexMetadata { + uuid: Uuid::new_v4(), + name: "cov_idx".to_string(), + fields: vec![vec_id, s_id], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.index.VectorIndexDetails".to_string(), + value: vec![], + })), + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![s_id], + }; + let create = |index: IndexMetadata| { + Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![index], + removed_indices: vec![], + }, + ) + }; + let rebase_over = |index: IndexMetadata| TransactionRebase { + transaction: create(index), + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::new(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + // Empty: the conflict must come from the schema comparison alone, not from + // any data-file path change. + covered_field_paths: Default::default(), + schema: Arc::new(before.clone()), + }; + let merge_with = |schema: lance_core::datatypes::Schema| { + Transaction::new_from_version( + 1, + Operation::Merge { + fragments: vec![], + schema, + preserves_nullability: true, + }, + ) + }; + + let mut rebase = rebase_over(covering_index.clone()); + assert!( + rebase.check_txn(&merge_with(after), 2).is_err(), + "a Merge that grows a covered struct's subtree must conflict" + ); + + // Control: a Merge that leaves the covered subtree alone keeps the optimistic + // path, so an ordinary add_columns still does not discard a covered build. + let mut rebase = rebase_over(covering_index); + assert!( + rebase.check_txn(&merge_with(before), 2).is_ok(), + "a Merge that does not touch the covered subtree must not conflict" + ); + } + async fn apply_deletion( delete_rows: &[u32], fragment: &mut Fragment, @@ -3483,6 +3954,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; for (other, expected_conflict) in other_transactions.iter().zip(expected_conflicts) { @@ -3687,6 +4160,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3746,6 +4221,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3887,6 +4364,8 @@ mod tests { affected_rows: affected_rows.as_ref(), conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let other_txn = Transaction::new(0, other.clone(), None); let result = rebase.check_txn(&other_txn, 1); @@ -3929,6 +4408,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = append_rebase.check_txn(&Transaction::new(0, merge.clone(), None), 1); assert_eq!( @@ -3944,6 +4425,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = merge_rebase.check_txn(&Transaction::new(0, append, None), 1); assert!( @@ -4020,6 +4503,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&Transaction::new(0, theirs, None), 1); assert_eq!( @@ -4154,6 +4639,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let update = Transaction::new( 1, @@ -4230,6 +4717,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&merge, 1); assert_eq!(result.is_err(), conflicts, "{result:?}"); @@ -4251,6 +4740,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: Arc::new(lance_core::datatypes::Schema::default()), }; let result = rebase.check_txn(&install, 1); assert!( @@ -4295,6 +4786,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let same_name = Transaction::new( @@ -4350,6 +4843,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let different_name_result = rebase.check_txn(&different_name, 1); assert!( @@ -4418,6 +4913,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&rewrite, 2); if expect_conflict { @@ -4765,6 +5262,124 @@ mod tests { assert!(rebase.check_txn(&txn2, 2).is_ok()); } + /// A covering-index build must treat its *covered* fields as dependencies during + /// conflict resolution: a concurrent write to a covered field would leave the + /// rebased index serving stale materialized values. + #[tokio::test] + async fn test_create_covering_index_conflicts_with_writes_to_covered_field() { + let dataset = test_dataset(10, 2).await; // fields: 0 = "a", 1 = "b" + + // `covering_fields` is always the trailing entries of `fields`. + let make_index = |name: &str, covering_fields: Vec| { + let mut fields = vec![0]; + fields.extend(covering_fields.iter().copied()); + IndexMetadata { + uuid: Uuid::new_v4(), + fields, + name: name.to_string(), + dataset_version: 1, + fragment_bitmap: Some([0, 1].into_iter().collect()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields, + } + }; + let create = |idx: IndexMetadata| { + Transaction::new_from_version( + 1, + Operation::CreateIndex { + new_indices: vec![idx], + removed_indices: vec![], + }, + ) + }; + let replace_field1 = || { + Transaction::new_from_version( + 1, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("replace.lance", vec![1], None), + )], + }, + ) + }; + // `fragment_field_paths` reads the last live file listing a field, which is how an + // in-place rewrite presents itself (the rewrite adds a file for the field and + // tombstones the old entry, and tombstones are skipped). So appending a file that + // carries `field` models a rewrite of it, and one carrying a brand-new id models + // `add_columns`. + let merge_appending_file_for = |field: i32, path: &str| { + let mut fragments = dataset.fragments().as_ref().clone(); + for fragment in fragments.iter_mut() { + fragment + .files + .push(DataFile::new_legacy_from_fields(path, vec![field], None)); + } + Transaction::new_from_version( + 1, + Operation::Merge { + fragments, + schema: dataset.schema().clone(), + preserves_nullability: true, + }, + ) + }; + let merge_rewriting_covered = || merge_appending_file_for(1, "rewritten.lance"); + let merge_adding_column = || merge_appending_file_for(2, "added.lance"); + + // Covering index (covers field 1): a replacement of field 1 conflicts... + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&replace_field1(), 2).is_err(), + "covering-index build must conflict with a data replacement of its covered field" + ); + // ...and so does a merge that rewrites the covered column's data file in place. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_rewriting_covered(), 2).is_err(), + "covering-index build must conflict with a merge that rewrites its covered field" + ); + // But a merge that only ADDS a column leaves the covered field's data untouched, + // so it must NOT discard the (potentially hours-long) covered index build. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("covering", vec![1])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_adding_column(), 2).is_ok(), + "covering-index build must not conflict with an add_columns that leaves covered data alone" + ); + + // Control: a plain index (no covered fields) keeps the optimistic path for + // both a merge and a replacement of the un-indexed field 1. + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("plain", vec![])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&merge_rewriting_covered(), 2).is_ok(), + "non-covering index build should keep the optimistic merge path" + ); + let mut rebase = + TransactionRebase::try_new(&dataset, create(make_index("plain", vec![])), None) + .await + .unwrap(); + assert!( + rebase.check_txn(&replace_field1(), 2).is_ok(), + "non-covering index build should not conflict with a replacement of an un-indexed field" + ); + } + /// Returns the IDs of fragments that have been modified by this operation. /// /// This does not include new fragments. @@ -5099,6 +5714,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&txn2, 1); @@ -5162,6 +5779,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5200,6 +5819,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5239,6 +5860,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5278,6 +5901,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5328,6 +5953,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result = rebase.check_txn(&committed_txn, 1); @@ -5353,6 +5980,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; let result_higher = rebase_higher.check_txn(&committed_txn, 1); @@ -5399,6 +6028,8 @@ mod tests { affected_rows: None, conflicting_frag_reuse_indices: Vec::new(), conflicting_mem_wal_compacted_sstables: Vec::new(), + covered_field_paths: Default::default(), + schema: empty_lance_schema(), }; // CreateIndex of MemWalIndex should be compatible with UpdateMemWalState diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 2bdc556abd9..75f97177de0 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -1443,6 +1443,21 @@ impl ANNIvfSubIndexExec { .first() .map(|idx| idx.covering_fields.clone()) .unwrap_or_default(); + // This exec declares one output schema but executes every delta segment, + // each emitting covering columns from its own storage -- so all deltas + // must agree on the covering set. The commit boundary enforces this; the + // check here defends against already-committed inconsistent metadata, + // failing with a clear error instead of a downstream schema mismatch. + if let Some(mismatch) = indices + .iter() + .find(|idx| idx.covering_fields != included_ids) + { + return Err(Error::index(format!( + "ANNIvfSubIndexExec: delta segments of index '{}' disagree on covering fields \ + ({:?} vs {:?}); the index metadata is inconsistent -- rebuild the index", + mismatch.name, included_ids, mismatch.covering_fields + ))); + } if !included_ids.is_empty() { let lance_schema = dataset.schema(); for id in &included_ids { @@ -1591,17 +1606,22 @@ impl ANNIvfEarlySearchResults { } } - fn record_batch(&self, batch: &RecordBatch) { + fn record_batch(&self, batch: &RecordBatch) -> DataFusionResult<()> { + let row_ids = batch.column_by_name(ROW_ID).ok_or_else(|| { + DataFusionError::Internal(format!( + "ANNIvfEarlySearchResults::record_batch: batch is missing the {ROW_ID} column" + )) + })?; let mut initial_ids = self.initial_ids.lock().unwrap(); let ids_to_record = (self.k - initial_ids.len()).min(batch.num_rows()); initial_ids.extend( - batch - .column(1) + row_ids .as_primitive::() .values() .iter() .take(ids_to_record), ); + Ok(()) } fn record_late_batch(&self, num_rows: usize) { @@ -1827,7 +1847,7 @@ impl ANNIvfSubIndexExec { } metrics.baseline_metrics.record_output(batch.num_rows()); if record_initial { - state.record_batch(&batch); + state.record_batch(&batch)?; } Ok(batch) }) @@ -2110,7 +2130,7 @@ impl ANNIvfSubIndexExec { seg_mask, ) .await?; - state.record_batch(&batch); + state.record_batch(&batch)?; Ok(batch) } }) @@ -2177,6 +2197,29 @@ impl ANNIvfSubIndexExec { // Fetch the covering columns (the fields after `[_distance, _rowid]`) from the // base table for the missing rows; the take is bounded by `not_found.len()`. + // Carry `_rowid` through as well: the take resolves ids with + // `MissingRowPolicy::Ignore` and silently returns FEWER rows than requested when an + // id no longer resolves through the row-id index -- a stale prefilter listing rows + // of a fragment a later delete removed entirely. Rebuilding the batch from + // `not_found` would then pair N distance/rowid values with < N covering values and + // fail with "all columns in a record batch must have the same length". Carrying + // `_rowid` lets the batch be rebuilt from the rows the take actually returned, + // which is what the non-covered path does when its `TakeExec` drops the same ids. + // + // A row that still resolves but is tombstoned would shorten the batch the same way. + // That is not reachable today, so the handling below is defensive for that case, not + // load-bearing. The invariant is enforced by `DatasetPreFilter`, not by the individual + // prefilter sources: its `deleted_ids` mask (`prefilter.rs`) is intersected into the final + // mask and covers every fragment in the manifest -- fragments inside the index bitmap + // contribute their deletion vectors, fragments outside it are blocked wholesale, and + // the mask is elided only when no fragment carries a deletion file at all. So a + // tombstoned address never enters `mask` regardless of which source produced it. + // Check that invariant there, not here, when adding a new `PreFilterSource`. + // + // `include_row_id()`, NOT `_rowid` in the projected column list: naming it there + // sets `must_add_row_offset`, which both arms the take's "must not target deleted + // rows" check (turning any shortfall into a hard error) and computes a row-offset + // column per batch that `project_batch` then discards. let covered_names: Vec = output_schema .fields() .iter() @@ -2190,17 +2233,40 @@ impl ANNIvfSubIndexExec { // identity (id == address) when it does not, so it is correct in both cases -- // unlike `try_new_from_addresses`, which would misread a stable id as // `frag = id >> 32, offset = id` and take the wrong physical row (or error). - let covered = TakeBuilder::try_new_from_ids(dataset, not_found.clone(), projection)? + let covered = TakeBuilder::try_new_from_ids(dataset, not_found, projection)? + .include_row_id() .execute() .await?; - let n = not_found.len(); - let mut columns: Vec = vec![ - Arc::new(Float32Array::from_value(f32::INFINITY, n)), - Arc::new(UInt64Array::from(not_found)), - ]; - columns.extend(covered.columns().iter().cloned()); - let batch = RecordBatch::try_new(output_schema, columns) + // Realign on what came back, not on what was asked for. + let returned_ids = covered + .column_by_name(ROW_ID) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "covered recovery take did not return the requested {ROW_ID} column" + )) + })? + .clone(); + let n = covered.num_rows(); + if n == 0 { + return Ok(None); + } + // Assemble by name. `include_row_id` appends `_rowid` to the take's output rather + // than placing it first, so no position in `covered` is contractual; the declared + // `[_distance, _rowid, ]` order comes from `output_schema` alone. + let mut columns: Vec = Vec::with_capacity(output_schema.fields().len()); + columns.push(Arc::new(Float32Array::from_value(f32::INFINITY, n))); + columns.push(returned_ids); + for field in output_schema.fields().iter().skip(2) { + let column = covered.column_by_name(field.name()).ok_or_else(|| { + DataFusionError::Internal(format!( + "covered recovery take did not return the requested column '{}'", + field.name() + )) + })?; + columns.push(column.clone()); + } + let batch = RecordBatch::try_new(output_schema.clone(), columns) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; Ok(Some(batch)) } @@ -2386,11 +2452,13 @@ impl ExecutionPlan for ANNIvfSubIndexExec { // Covered null-vector recovery: track which row ids the search emits, then after // it finishes emit any prefilter-matched rows that had no index entry, with their // covering columns taken from the base table (see `covered_not_found_batch`). - // Tracking is skipped when recovery provably cannot emit anything: without a - // prefilter source there is no allow-list, so `covered_not_found_batch` always - // returns `None` -- the per-batch lock and inserts would be pure overhead on - // the covered hot path. - let track_emitted = has_covered && !matches!(self.prefilter_source, PreFilterSource::None); + // This must NOT be gated on `prefilter_source`: `DatasetPreFilter` builds a bounded + // allow-list from the deletion mask alone, so an unfiltered query over a dataset with + // deletions (or a fragment outside the index bitmap) still arms recovery. Skipping the + // bookkeeping there left `emitted` empty, and recovery re-emitted every live row on top + // of the ones the search had already returned. The cost is already confined to covered + // plans, and recovery itself stays gated behind `has_covered`. + let track_emitted = has_covered; let emitted_row_ids = Arc::new(std::sync::Mutex::new(roaring::RoaringTreemap::new())); let emitted_for_inspect = emitted_row_ids.clone(); let recon_prefilter = pre_filter.clone(); @@ -2398,6 +2466,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let recon_ds = ds.clone(); let recon_schema = schema.clone(); let recon_k = query.k; + let declared_schema = schema.clone(); Ok(Box::pin(RecordBatchStreamAdapter::new( schema, @@ -2483,6 +2552,30 @@ impl ExecutionPlan for ANNIvfSubIndexExec { // Each delta stream is split into an early and late search. The late search // will not start until the early search is complete across all deltas. .try_flatten_unordered(None) + // The search emits covering columns read from index STORAGE, while this node + // declares them from the manifest (`IndexMetadata.covering_fields`). The two can + // disagree: `FLAG_COVERED_INDEX_METADATA` is a best-effort writer fence, so a + // client predating it can drop the declaration while the index files keep the + // payload. Without narrowing, this stream mixes batch widths -- a wider search + // batch next to the narrower non-covered shortcut batch -- which panics inside + // DataFusion's TopK (`interleave_record_batch` reads every batch through the first + // batch's schema). Only ever narrows: the opposite desync (declared but absent from + // storage) is prevented at build time (covering requires the V3 index format and + // rejects columns the storage cannot emit) and otherwise fails loudly, which is the + // right outcome. + // + // Caveat worth knowing: the test below is cardinality-based, so it cannot detect a + // desync in which storage holds the same *number* of covering columns as the + // declaration but different ones. + .map(move |batch| { + let batch = batch?; + if batch.num_columns() > declared_schema.fields().len() { + return batch + .project_by_schema(declared_schema.as_ref()) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)); + } + Ok(batch) + }) .inspect_ok(move |batch| { // Record emitted row ids so the reconciliation below can find the // prefilter rows the search never returned (no index entry). @@ -2945,9 +3038,14 @@ mod tests { .unwrap(), ); + // A legal covering shape (`covering_fields` is the trailing slice of `fields`, + // which keeps at least one keyed entry -- see + // `IndexMetadata::validate_covering_fields`) whose carried id 9999 the schema + // cannot resolve. Keeping the shape legal is what makes this a test of the + // resolution failure rather than of metadata no commit can produce. let index = IndexMetadata { uuid: uuid::Uuid::new_v4(), - fields: vec![], + fields: vec![0, 9999], name: "vector_idx".to_string(), dataset_version: 1, fragment_bitmap: Some(RoaringBitmap::new()), @@ -2958,6 +3056,9 @@ mod tests { files: None, covering_fields: vec![9999], }; + index + .validate_covering_fields() + .expect("the fixture must be a legal covering declaration"); let result = ANNIvfSubIndexExec::try_new( input, @@ -2973,6 +3074,77 @@ mod tests { ); } + /// The exec declares one output schema but executes every delta segment, so delta + /// segments that disagree on covering fields must be rejected up front with a clear + /// error, not fail later with a stream/plan schema mismatch. (The commit boundary + /// rejects new mixed sets; this defends against already-committed metadata.) + #[tokio::test] + async fn test_ann_sub_index_rejects_mixed_covering_deltas() { + use lance_datafusion::exec::OneShotExec; + + let input = Arc::new(OneShotExec::from_batch(RecordBatch::new_empty( + KNN_PARTITION_SCHEMA.clone(), + ))); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-mixed-covering-deltas", + None, + ) + .await + .unwrap(), + ); + let id_field_id = dataset.schema().field("id").unwrap().id; + // This exec path only resolves `covering_fields` against the live schema, so a + // key id that is not itself a real schema field is harmless here; it stands in + // for the vector key column a real covering index would also carry in `fields`. + let key_field_id = id_field_id + 1; + + // `covering_fields` is always the trailing entries of `fields`. + let covered_delta = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![key_field_id, id_field_id], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + covering_fields: vec![id_field_id], + }; + let plain_delta = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + covering_fields: Vec::new(), + ..covered_delta.clone() + }; + + let result = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![covered_delta, plain_delta], + base_query(), + PreFilterSource::None, + ); + let err = result.expect_err("delta segments with mixed covering must be rejected"); + assert!( + err.to_string().contains("disagree on covering fields"), + "expected a delta covering mismatch error, got: {err}" + ); + } + #[test] fn test_effective_query_parallelism_clamps_to_cpu_pool() { let mut query = base_query(); @@ -3825,16 +3997,18 @@ mod tests { query.minimum_nprobes = 0; query.maximum_nprobes = Some(3); let state = Arc::new(ANNIvfEarlySearchResults::new(1, query.k)); - state.record_batch( - &RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(vec![0.0])), - Arc::new(UInt64Array::from(vec![999])), - ], + state + .record_batch( + &RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(vec![0.0])), + Arc::new(UInt64Array::from(vec![999])), + ], + ) + .unwrap(), ) - .unwrap(), - ); + .unwrap(); let prefilter = empty_prefilter().await; let batches = ANNIvfSubIndexExec::late_search( @@ -3860,6 +4034,32 @@ mod tests { assert_eq!(state.num_results_found.load(Ordering::Relaxed), 2); } + /// `record_batch` must resolve `_rowid` by name, not by position: a covered ANN + /// batch appends covering columns after `_rowid`, but nothing here should assume + /// that specific layout. Puts another UInt64 column ahead of `_rowid` so a + /// positional read would silently record the decoy column's values instead. + #[test] + fn test_early_search_results_record_batch_resolves_row_id_by_name() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new(DIST_COL, DataType::Float32, true), + ArrowField::new("decoy_u64_col", DataType::UInt64, true), + ROW_ID_FIELD.clone(), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Float32Array::from(vec![0.0, 1.0])), + Arc::new(UInt64Array::from(vec![111, 222])), + Arc::new(UInt64Array::from(vec![7, 8])), + ], + ) + .unwrap(); + + let state = ANNIvfEarlySearchResults::new(1, 2); + state.record_batch(&batch).unwrap(); + assert_eq!(*state.initial_ids.lock().unwrap(), vec![7, 8]); + } + fn row_ids_of(batches: &[RecordBatch]) -> Vec { batches .iter() From f0fcf3079a22109b9313e45ce7e3bef2b4ced7cd Mon Sep 17 00:00:00 2001 From: Vivek Date: Mon, 17 Aug 2026 15:27:39 -0700 Subject: [PATCH 4/4] feat(index)!: support covering columns in distributed vector index builds A vector index built as separate shards and merged afterwards now carries its covering columns through instead of rejecting them. Each shard stores its covered values and the merge step classifies them from the shard schema, using each storage format's own list of internal column names so a new internal column is excluded everywhere at once. Both distributed commit styles are covered by tests that check query results against the base table. BREAKING CHANGE: the four public `init_writer_for_{flat,pq,sq,rq}` functions in lance-index each take a new `covering_fields` argument. Pass an empty slice for the previous behaviour. `VectorStore` also gains a required `INTERNAL_COLUMNS` associated const naming the storage's own non-covering columns. No compatibility overloads were added: an empty slice is trivially expressible at the call site, and carrying a second name for each function would outlive the reason for it. --- rust/lance-index/src/vector/bq/storage.rs | 46 +- .../src/vector/distributed/index_merger.rs | 224 +++- rust/lance-index/src/vector/flat/storage.rs | 115 +-- rust/lance-index/src/vector/pq/storage.rs | 19 +- rust/lance-index/src/vector/sq/storage.rs | 23 +- rust/lance-index/src/vector/storage.rs | 33 +- rust/lance/src/index.rs | 256 ++++- rust/lance/src/index/vector.rs | 344 ++++--- rust/lance/src/index/vector/builder.rs | 139 ++- rust/lance/src/index/vector/ivf/v2.rs | 968 +++++++++++++++++- 10 files changed, 1799 insertions(+), 368 deletions(-) diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index cf719789e94..4e4d10a12f2 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -67,7 +67,6 @@ use crate::vector::pq::storage::transpose; use crate::vector::quantizer::{QuantizerMetadata, QuantizerStorage}; use crate::vector::storage::{ DistCalculator, DistanceCalculatorOptions, QueryResidual, RabitRawQueryContext, VectorStore, - covering_field_indices_excluding, }; pub const RABIT_METADATA_KEY: &str = "lance:rabit"; @@ -80,6 +79,25 @@ pub const RABIT_EX_CODE_COLUMN: &str = "__ex_codes"; /// older versions, which fail with a missing-column error instead of /// misinterpreting the bytes. pub const RABIT_BLOCKED_EX_CODE_COLUMN: &str = "__blocked_ex_codes"; +/// RaBitQ storage's internal (non-covering) column names: the row id, the binary code +/// and extended-bit code columns, the per-row factor columns (some present only for +/// higher `num_bits`), and the legacy `__ivf_part_id` column (left in pre-#3606 +/// single-partition builds; unlike PQ, RaBitQ storage does not drop it at +/// construction). Every other column is a user-declared covering ("included") column. +/// The single authority for this set -- the distributed shard merger classifies from +/// it too, so a new factor column added here is excluded everywhere at once. +pub const RABIT_INTERNAL_COLUMNS: &[&str] = &[ + ROW_ID, + RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + PART_ID_COLUMN, +]; pub const SEGMENT_LENGTH: usize = 4; pub const SEGMENT_NUM_CODES: usize = 1 << SEGMENT_LENGTH; const RABIT_PRUNE_STATS_ENV: &str = "LANCE_RQ_PRUNE_STATS"; @@ -2042,6 +2060,9 @@ fn accumulate_filtered_distances_into_heap( impl VectorStore for RabitQuantizationStorage { type DistanceCalculator<'a> = RabitDistCalculator<'a>; + /// RaBitQ storage is `[_rowid, , ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = RABIT_INTERNAL_COLUMNS; + fn as_any(&self) -> &dyn std::any::Any { self } @@ -2050,29 +2071,6 @@ impl VectorStore for RabitQuantizationStorage { self.batch.schema_ref() } - /// RQ storage carries the row id plus many internal columns (the binary codes, the - /// extended-bit codes, and the per-row factor columns, some present only for higher - /// `num_bits`); the covering ("included") columns are everything else. - fn covering_field_indices(&self) -> Vec { - // Every RaBitQ-internal column plus the legacy `__ivf_part_id` column - // (left in pre-#3606 single-partition builds; unlike PQ, RaBitQ storage - // does not drop it at construction, so it is excluded here). Anything - // else is a user-declared covering column. - const INTERNAL: &[&str] = &[ - ROW_ID, - RABIT_CODE_COLUMN, - RABIT_EX_CODE_COLUMN, - RABIT_BLOCKED_EX_CODE_COLUMN, - ADD_FACTORS_COLUMN, - SCALE_FACTORS_COLUMN, - ERROR_FACTORS_COLUMN, - EX_ADD_FACTORS_COLUMN, - EX_SCALE_FACTORS_COLUMN, - PART_ID_COLUMN, - ]; - covering_field_indices_excluding(self.schema().as_ref(), INTERNAL) - } - fn to_batches(&self) -> Result + Send> { Ok(std::iter::once(self.batch.clone())) } diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 370011eb8bd..cc72ef43957 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -21,8 +21,8 @@ use crate::IndexMetadata as IndexMetaSchema; use crate::pb; use crate::scalar::OldIndexDataFilter; use crate::vector::bq::storage::{ - RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RabitQuantizationMetadata, RabitQueryEstimator, - pack_codes, rabit_binary_code_field, rabit_ex_code_field, + RABIT_CODE_COLUMN, RABIT_INTERNAL_COLUMNS, RABIT_METADATA_KEY, RabitQuantizationMetadata, + RabitQueryEstimator, pack_codes, rabit_binary_code_field, rabit_ex_code_field, }; use crate::vector::bq::transform::{ ADD_FACTORS_FIELD, ERROR_FACTORS_FIELD, EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD, @@ -30,14 +30,20 @@ use crate::vector::bq::transform::{ }; use crate::vector::bq::validate_rq_num_bits; use crate::vector::flat::index::FlatMetadata; +use crate::vector::flat::storage::{FLAT_COLUMN, FLAT_INTERNAL_COLUMNS}; use crate::vector::ivf::storage::{IVF_METADATA_KEY, IvfModel as IvfStorageModel}; -use crate::vector::pq::storage::{PQ_METADATA_KEY, ProductQuantizationMetadata, transpose}; +use crate::vector::pq::storage::{ + PQ_INTERNAL_COLUMNS, PQ_METADATA_KEY, ProductQuantizationMetadata, transpose, +}; use crate::vector::quantizer::QuantizerMetadata; -use crate::vector::sq::storage::{SQ_METADATA_KEY, ScalarQuantizationMetadata}; +use crate::vector::sq::storage::{ + SQ_INTERNAL_COLUMNS, SQ_METADATA_KEY, ScalarQuantizationMetadata, +}; use crate::vector::storage::STORAGE_METADATA_KEY; +use crate::vector::storage::{COVERING_FIELD_IDS_KEY, covering_field_indices_excluding}; use crate::vector::{DISTANCE_TYPE_KEY, PQ_CODE_COLUMN, SQ_CODE_COLUMN}; use crate::{INDEX_AUXILIARY_FILE_NAME, INDEX_METADATA_SCHEMA_KEY}; -use arrow_schema::{DataType, Field, Schema as ArrowSchema}; +use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema}; use bytes::Bytes; use lance_core::datatypes::Schema as LanceSchema; use lance_file::reader::{FileReader as V2Reader, FileReaderOptions as V2ReaderOptions}; @@ -249,6 +255,34 @@ fn init_writer_for_storage( Ok(()) } +/// The covering ("included") columns carried in a shard's auxiliary storage schema: every +/// field that is not the row id or one of the quantizer's internal code/factor columns (or the +/// legacy `__ivf_part_id`). The merger rebuilds the unified output schema from quantizer +/// metadata (it reshapes the code column via the transpose), so these fields must be +/// re-appended or the covered payload is silently dropped from the merged index. All shards +/// share the same covering set, so the first shard's schema is authoritative. +fn covering_fields_from_shard_schema( + schema: &ArrowSchema, + index_type: SupportedIvfIndexType, +) -> Vec { + use SupportedIvfIndexType::*; + // Each storage's exported internal set already excludes PART_ID_COLUMN, which + // matters here: a legacy PQ shard's raw aux schema may still physically carry + // `__ivf_part_id` even though PQ storage drops it at load, so classifying from + // the raw schema (pre-load) must exclude it or it would be wrongly classified + // as a covering column. + let internal: &[&str] = match index_type { + IvfPq | IvfHnswPq => PQ_INTERNAL_COLUMNS, + IvfSq | IvfHnswSq => SQ_INTERNAL_COLUMNS, + IvfFlat | IvfHnswFlat => FLAT_INTERNAL_COLUMNS, + IvfRq => RABIT_INTERNAL_COLUMNS, + }; + covering_field_indices_excluding(schema, internal) + .into_iter() + .map(|i| schema.fields()[i].clone()) + .collect() +} + /// Create and initialize a unified writer for FLAT storage. pub async fn init_writer_for_flat( object_store: &lance_io::object_store::ObjectStore, @@ -257,18 +291,21 @@ pub async fn init_writer_for_flat( item_type: &DataType, dt: DistanceType, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( - crate::vector::flat::storage::FLAT_COLUMN, + FLAT_COLUMN, DataType::FixedSizeList( Arc::new(Field::new("item", item_type.clone(), true)), d0 as i32, ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -291,13 +328,14 @@ pub async fn init_writer_for_pq( dt: DistanceType, pm: &ProductQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let num_bytes = if pm.nbits == 4 { pm.num_sub_vectors / 2 } else { pm.num_sub_vectors }; - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( PQ_CODE_COLUMN, @@ -307,7 +345,9 @@ pub async fn init_writer_for_pq( ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -336,9 +376,10 @@ pub async fn init_writer_for_sq( dt: DistanceType, sq_meta: &ScalarQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let d0 = sq_meta.dim; - let arrow_schema = ArrowSchema::new(vec![ + let mut fields = vec![ (*ROW_ID_FIELD).clone(), Field::new( SQ_CODE_COLUMN, @@ -348,7 +389,9 @@ pub async fn init_writer_for_sq( ), true, ), - ]); + ]; + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); + let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( format_version, @@ -368,6 +411,7 @@ pub async fn init_writer_for_rq( dt: DistanceType, rq_meta: &RabitQuantizationMetadata, format_version: ConcreteFileVersion, + covering_fields: &[FieldRef], ) -> Result { let mut fields = vec![ (*ROW_ID_FIELD).clone(), @@ -383,6 +427,7 @@ pub async fn init_writer_for_rq( fields.push(EX_ADD_FACTORS_FIELD.clone()); fields.push(EX_SCALE_FACTORS_FIELD.clone()); } + fields.extend(covering_fields.iter().map(|f| f.as_ref().clone())); let arrow_schema = ArrowSchema::new(fields); let writer = object_store.create(aux_out).await?; let mut w = versions::create_writer( @@ -898,6 +943,11 @@ async fn merge_partial_vector_auxiliary_files_inner( let mut nlist_opt: Option = None; let mut accumulated_lengths: Vec = Vec::new(); let mut first_centroids: Option = None; + // Covering ("included") columns of the first shard, used as the authoritative set every + // later shard must match (see the per-shard check below). + // (covering fields, source dataset field ids) of the first shard; every later + // shard must match both. + let mut first_covering_fields: Option<(Vec, Option)> = None; // Track per-shard readers, IVF lengths, and precomputed partition offsets. // This avoids reopening each shard file for every partition during merge. @@ -1044,6 +1094,68 @@ async fn merge_partial_vector_auxiliary_files_inner( // Preserve the historical fallback while keeping the writer boundary exact. let fv = format_version.unwrap_or(ConcreteFileVersion::V2_0); + // Covering ("included") columns carried in this shard's storage. The unified writer + // schemas below are rebuilt from quantizer metadata, so these must be re-appended or + // the covered payload is dropped from the merged index. The output schema is built + // from the *first* shard's covering fields and `concat_batches` combines columns + // positionally, so a later shard whose covering columns differ in name, type, or order + // would be silently stored under the first shard's names -> wrong covered values. + // Reject any such mismatch rather than corrupt the merged index. + let shard_arrow: ArrowSchema = reader.schema().as_ref().into(); + let covering_fields = covering_fields_from_shard_schema(&shard_arrow, idx_type); + // Name and type alone cannot establish that two shards cover the *same logical + // column*: Arrow fields carry no Lance field id, and the ids in a shard's own + // Lance schema are file-local. A column dropped and re-added between two shard + // builds yields the same name and type under a different field id, and + // `concat_batches` would stack the old values under the new field. The build + // stamps the source field ids for exactly this comparison. + let covering_source_ids = reader + .metadata() + .file_schema + .metadata + .get(COVERING_FIELD_IDS_KEY) + .cloned(); + if !covering_fields.is_empty() && covering_source_ids.is_none() { + return Err(Error::index(format!( + "Distributed merge: shard {idx} carries covering (included) columns but records \ + no source field ids, so it cannot be proven to cover the same fields as the \ + other shards. Rebuild the shard." + ))); + } + match first_covering_fields.as_ref() { + None => { + first_covering_fields = Some((covering_fields.clone(), covering_source_ids.clone())) + } + Some((first, first_source_ids)) => { + let matches = first.len() == covering_fields.len() + && first + .iter() + .zip(&covering_fields) + .all(|(a, b)| a.name() == b.name() && a.data_type() == b.data_type()) + && *first_source_ids == covering_source_ids; + if !matches { + let describe = |fields: &[FieldRef], source_ids: &Option| { + let names = fields + .iter() + .map(|f| format!("{}: {}", f.name(), f.data_type())) + .collect::>() + .join(", "); + match source_ids { + Some(ids) => format!("{names} (source field ids {ids})"), + None => names, + } + }; + return Err(Error::index(format!( + "Distributed merge: covering (included) columns differ across shards; \ + shard 0 has [{}] but shard {idx} has [{}]. All shards must declare the \ + same covering columns (same source fields, names, types, and order).", + describe(first, first_source_ids), + describe(&covering_fields, &covering_source_ids), + ))); + } + } + } + match idx_type { SupportedIvfIndexType::IvfSq => { // Handle Scalar Quantization (SQ) storage for IVF_SQ @@ -1097,8 +1209,15 @@ async fn merge_partial_vector_auxiliary_files_inner( sq_meta = Some(sq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_sq(object_store, &aux_out, dt, &sq_meta_parsed, fv).await?; + let w = init_writer_for_sq( + object_store, + &aux_out, + dt, + &sq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1205,8 +1324,15 @@ async fn merge_partial_vector_auxiliary_files_inner( rq_meta = Some(rq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_rq(object_store, &aux_out, dt, &rq_meta_parsed, fv).await?; + let w = init_writer_for_rq( + object_store, + &aux_out, + dt, + &rq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1305,8 +1431,15 @@ async fn merge_partial_vector_auxiliary_files_inner( if v2w_opt.is_none() { let mut pm_for_unified = pm.clone(); pm_for_unified.transposed = true; - let w = - init_writer_for_pq(object_store, &aux_out, dt, &pm_for_unified, fv).await?; + let w = init_writer_for_pq( + object_store, + &aux_out, + dt, + &pm_for_unified, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1334,8 +1467,16 @@ async fn merge_partial_vector_auxiliary_files_inner( return Err(Error::index("Dimension mismatch across shards".to_string())); } if v2w_opt.is_none() { - let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv) - .await?; + let w = init_writer_for_flat( + object_store, + &aux_out, + d0, + &item_type, + dt, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1367,8 +1508,16 @@ async fn merge_partial_vector_auxiliary_files_inner( return Err(Error::index("Dimension mismatch across shards".to_string())); } if v2w_opt.is_none() { - let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv) - .await?; + let w = init_writer_for_flat( + object_store, + &aux_out, + d0, + &item_type, + dt, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1463,8 +1612,15 @@ async fn merge_partial_vector_auxiliary_files_inner( if v2w_opt.is_none() { let mut pm_for_unified = pm.clone(); pm_for_unified.transposed = true; - let w = - init_writer_for_pq(object_store, &aux_out, dt, &pm_for_unified, fv).await?; + let w = init_writer_for_pq( + object_store, + &aux_out, + dt, + &pm_for_unified, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1515,8 +1671,15 @@ async fn merge_partial_vector_auxiliary_files_inner( sq_meta = Some(sq_meta_parsed.clone()); } if v2w_opt.is_none() { - let w = - init_writer_for_sq(object_store, &aux_out, dt, &sq_meta_parsed, fv).await?; + let w = init_writer_for_sq( + object_store, + &aux_out, + dt, + &sq_meta_parsed, + fv, + &covering_fields, + ) + .await?; v2w_opt = Some(w); } } @@ -1716,6 +1879,15 @@ async fn merge_partial_vector_auxiliary_files_inner( for len in merged_lengths.iter() { ivf_model.add_partition(*len); } + // Carry the covering provenance onto the merged output. Every shard has already + // been checked to agree on this stamp, so the first shard's value IS the merged + // file's value. Without it the merged segment carries covering columns but no + // source ids, and feeding it back in as a shard of a later hierarchical merge + // would be rejected by the very check that verified its inputs -- with an error + // telling the user to rebuild a shard the merger itself produced. + if let Some((_, Some(covering_source_ids))) = first_covering_fields.as_ref() { + w.add_schema_metadata(COVERING_FIELD_IDS_KEY, covering_source_ids.clone()); + } let dt2 = distance_type.ok_or_else(|| Error::index("Distance type missing".to_string()))?; write_unified_ivf_and_index_metadata(w, &ivf_model, dt2, idx_type_final).await?; let summary = w.finish().await?; diff --git a/rust/lance-index/src/vector/flat/storage.rs b/rust/lance-index/src/vector/flat/storage.rs index 3c90f0c9e01..de84897e92e 100644 --- a/rust/lance-index/src/vector/flat/storage.rs +++ b/rust/lance-index/src/vector/flat/storage.rs @@ -8,9 +8,7 @@ use crate::frag_reuse::{FragReuseIndex, FragReuseIndexHandle}; use crate::scalar::RowIdRemapper; use crate::vector::PART_ID_COLUMN; use crate::vector::quantizer::QuantizerStorage; -use crate::vector::storage::{ - DistCalculator, VectorStore, covering_field_indices_excluding, remap_row_ids_by_name, -}; +use crate::vector::storage::{DistCalculator, VectorStore, remap_row_ids_by_name}; use crate::vector::utils::do_prefetch; use arrow::array::AsArray; use arrow::compute::concat_batches; @@ -29,6 +27,14 @@ use lance_linalg::distance::{Cosine, DistanceType, Dot, L2, Normalize, norm_l2_f pub const FLAT_COLUMN: &str = "flat"; +/// Flat storage's internal (non-covering) column names: the row id, the flat vector +/// column, and the legacy `__ivf_part_id` column (left in pre-#3606 single-partition +/// builds; unlike PQ, flat storage does not drop it at construction). Every other +/// column in a flat storage schema is a user-declared covering ("included") column. +/// The single authority for this set -- the distributed shard merger classifies from +/// it too, so a new internal column added here is excluded everywhere at once. +pub const FLAT_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN]; + /// Per-vector L2 norms cached for Cosine distance, so `cosine_with_norms` can /// skip recomputing each stored vector's norm per comparison. `None` for other /// metrics and for value types without a norm kernel. @@ -166,6 +172,9 @@ impl FlatFloatStorage { impl VectorStore for FlatFloatStorage { type DistanceCalculator<'a> = FlatFloatDistanceCalc<'a>; + /// Flat storage is `[_rowid, flat, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = FLAT_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok([self.batch.clone()].into_iter()) } @@ -197,17 +206,6 @@ impl VectorStore for FlatFloatStorage { self.batch.schema_ref() } - /// Flat storage is `[_rowid, flat, ]`, so the covering columns - /// are every field except the row id, the flat vector column, and the legacy - /// `__ivf_part_id` column (left in pre-#3606 single-partition builds; unlike - /// PQ, flat storage does not drop it at construction, so it is excluded here). - fn covering_field_indices(&self) -> Vec { - covering_field_indices_excluding( - self.schema().as_ref(), - &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN], - ) - } - fn as_any(&self) -> &dyn std::any::Any { self } @@ -353,6 +351,9 @@ impl FlatBinStorage { impl VectorStore for FlatBinStorage { type DistanceCalculator<'a> = FlatDistanceCal<'a, UInt8Type>; + /// Flat storage is `[_rowid, flat, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = FLAT_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok([self.batch.clone()].into_iter()) } @@ -383,17 +384,6 @@ impl VectorStore for FlatBinStorage { self.batch.schema_ref() } - /// Flat storage is `[_rowid, flat, ]`, so the covering columns - /// are every field except the row id, the flat vector column, and the legacy - /// `__ivf_part_id` column (left in pre-#3606 single-partition builds; unlike - /// PQ, flat storage does not drop it at construction, so it is excluded here). - fn covering_field_indices(&self) -> Vec { - covering_field_indices_excluding( - self.schema().as_ref(), - &[ROW_ID, FLAT_COLUMN, PART_ID_COLUMN], - ) - } - fn as_any(&self) -> &dyn std::any::Any { self } @@ -716,60 +706,61 @@ mod tests { // `IndexMetadata.covering_fields`, empty for those indexes) lacks, // desyncing the two and failing every query on that legacy index. #[test] - fn test_flat_float_covering_excludes_legacy_part_id_column() { + fn test_flat_covering_excludes_legacy_part_id_column() { use crate::vector::PART_ID_COLUMN; - const DIM: i32 = 4; + const DIM: i32 = 8; const N: usize = 8; - let row_ids = UInt64Array::from_iter_values(0..N as u64); - let values = + + let part_id_batch = |vectors: ArrayRef| { + RecordBatch::try_from_iter_with_nullable(vec![ + ( + ROW_ID, + Arc::new(UInt64Array::from_iter_values(0..N as u64)) as ArrayRef, + true, + ), + (FLAT_COLUMN, vectors, true), + ( + PART_ID_COLUMN, + Arc::new(arrow_array::UInt32Array::from_iter_values( + (0..N).map(|_| 0), + )) as ArrayRef, + true, + ), + ]) + .unwrap() + }; + let meta = FlatMetadata { dim: DIM as usize }; + + let floats = arrow_array::Float32Array::from_iter_values((0..N * DIM as usize).map(|v| v as f32)); - let vectors = FixedSizeListArray::try_new_from_values(values, DIM).unwrap(); - let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0)); - let batch = RecordBatch::try_from_iter_with_nullable(vec![ - (ROW_ID, Arc::new(row_ids) as ArrayRef, true), - (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true), - (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true), - ]) - .unwrap(); - let storage = FlatFloatStorage::try_from_batch( - batch, - &FlatMetadata { dim: DIM as usize }, + let float_storage = FlatFloatStorage::try_from_batch( + part_id_batch(Arc::new( + FixedSizeListArray::try_new_from_values(floats, DIM).unwrap(), + )), + &meta, DistanceType::L2, None, ) .unwrap(); assert!( - storage.covering_field_indices().is_empty(), - "legacy __ivf_part_id must not be classified as a covering column" + float_storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as covering (float storage)" ); - } - #[test] - fn test_flat_bin_covering_excludes_legacy_part_id_column() { - use crate::vector::PART_ID_COLUMN; - const DIM: i32 = 8; - const N: usize = 8; - let row_ids = UInt64Array::from_iter_values(0..N as u64); let codes = arrow_array::UInt8Array::from_iter_values((0..N * DIM as usize).map(|v| v as u8)); - let vectors = FixedSizeListArray::try_new_from_values(codes, DIM).unwrap(); - let part_ids = arrow_array::UInt32Array::from_iter_values((0..N).map(|_| 0)); - let batch = RecordBatch::try_from_iter_with_nullable(vec![ - (ROW_ID, Arc::new(row_ids) as ArrayRef, true), - (FLAT_COLUMN, Arc::new(vectors) as ArrayRef, true), - (PART_ID_COLUMN, Arc::new(part_ids) as ArrayRef, true), - ]) - .unwrap(); - let storage = FlatBinStorage::try_from_batch( - batch, - &FlatMetadata { dim: DIM as usize }, + let bin_storage = FlatBinStorage::try_from_batch( + part_id_batch(Arc::new( + FixedSizeListArray::try_new_from_values(codes, DIM).unwrap(), + )), + &meta, DistanceType::L2, None, ) .unwrap(); assert!( - storage.covering_field_indices().is_empty(), - "legacy __ivf_part_id must not be classified as a covering column" + bin_storage.covering_field_indices().is_empty(), + "legacy __ivf_part_id must not be classified as covering (binary storage)" ); } diff --git a/rust/lance-index/src/vector/pq/storage.rs b/rust/lance-index/src/vector/pq/storage.rs index 0bcd9a3a5ef..c2d7c0f340d 100644 --- a/rust/lance-index/src/vector/pq/storage.rs +++ b/rust/lance-index/src/vector/pq/storage.rs @@ -46,13 +46,21 @@ use crate::{ PART_ID_COLUMN, PQ_CODE_COLUMN, pq::transform::PQTransformer, quantizer::{QuantizerMetadata, QuantizerStorage}, - storage::{DistCalculator, VectorStore, covering_field_indices_excluding}, + storage::{DistCalculator, VectorStore}, transform::Transformer, }, }; pub const PQ_METADATA_KEY: &str = "lance:pq"; +/// PQ storage's internal (non-covering) column names: the row id, the PQ code column, +/// and the legacy `__ivf_part_id` column (written by pre-#3606 single-partition +/// builds; PQ storage drops it at construction, so post-load it is never present and +/// excluding it is a no-op -- but pre-load classification, e.g. from a raw shard +/// schema in the distributed merger, needs it). Every other column is a user-declared +/// covering ("included") column. The single authority for this set. +pub const PQ_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, PQ_CODE_COLUMN, PART_ID_COLUMN]; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProductQuantizationMetadata { pub codebook_position: usize, @@ -770,6 +778,9 @@ impl QuantizerStorage for ProductQuantizationStorage { impl VectorStore for ProductQuantizationStorage { type DistanceCalculator<'a> = PQDistCalculator; + /// PQ storage is `[_rowid, __pq_code, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = PQ_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok(std::iter::once(self.batch.clone())) } @@ -782,12 +793,6 @@ impl VectorStore for ProductQuantizationStorage { self.batch.schema_ref() } - /// PQ storage is `[_rowid, __pq_code, ]`, so the covering - /// columns are every field except the row id and the PQ code column. - fn covering_field_indices(&self) -> Vec { - covering_field_indices_excluding(self.schema().as_ref(), &[ROW_ID, PQ_CODE_COLUMN]) - } - fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/rust/lance-index/src/vector/sq/storage.rs b/rust/lance-index/src/vector/sq/storage.rs index c7b42c4c33c..35fd8a64acf 100644 --- a/rust/lance-index/src/vector/sq/storage.rs +++ b/rust/lance-index/src/vector/sq/storage.rs @@ -34,7 +34,7 @@ use crate::{ quantizer::{QuantizerMetadata, QuantizerStorage}, storage::{ DistCalculator, DistanceCalculatorOptions, QueryResidual, VectorStore, - covering_field_indices_excluding, remap_row_ids_by_name, + remap_row_ids_by_name, }, transform::Transformer, }, @@ -42,6 +42,13 @@ use crate::{ pub const SQ_METADATA_KEY: &str = "lance:sq"; +/// SQ storage's internal (non-covering) column names: the row id, the SQ code column, +/// and the legacy `__ivf_part_id` column (left in pre-#3606 single-partition builds; +/// unlike PQ, SQ does not drop it at construction). Every other column in an SQ +/// storage schema is a user-declared covering ("included") column. The single +/// authority for this set -- the distributed shard merger classifies from it too. +pub const SQ_INTERNAL_COLUMNS: &[&str] = &[ROW_ID, SQ_CODE_COLUMN, PART_ID_COLUMN]; + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ScalarQuantizationMetadata { pub dim: usize, @@ -345,6 +352,9 @@ impl QuantizerStorage for ScalarQuantizationStorage { impl VectorStore for ScalarQuantizationStorage { type DistanceCalculator<'a> = SQDistCalculator<'a>; + /// SQ storage is `[_rowid, __sq_code, ]`. + const INTERNAL_COLUMNS: &'static [&'static str] = SQ_INTERNAL_COLUMNS; + fn to_batches(&self) -> Result> { Ok(self.chunks.iter().map(|c| c.batch.clone())) } @@ -373,17 +383,6 @@ impl VectorStore for ScalarQuantizationStorage { self.chunks[0].schema() } - /// SQ storage is `[_rowid, __sq_code, ]`, so the covering - /// columns are every field except the row id, the SQ code column, and the - /// legacy `__ivf_part_id` column (left in pre-#3606 single-partition builds; - /// unlike PQ, SQ does not drop it at construction, so it is excluded here). - fn covering_field_indices(&self) -> Vec { - covering_field_indices_excluding( - self.schema().as_ref(), - &[ROW_ID, SQ_CODE_COLUMN, PART_ID_COLUMN], - ) - } - fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 9dc843b9d6f..493e28d200d 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -155,6 +155,19 @@ pub trait DistCalculator { pub const STORAGE_METADATA_KEY: &str = "storage_metadata"; +/// Schema-metadata key recording the *source dataset* field ids of a storage file's +/// covering ("included") columns, comma separated in declaration order. +/// +/// Arrow fields carry no Lance field id, so a storage file's own schema cannot say +/// which logical column a covering column came from -- and the ids in the file's Lance +/// schema are file-local, assigned when the writer was built. Without this, a +/// distributed merge comparing shards by name and type would accept two shards whose +/// covering columns share a name and type but belong to different fields (a column +/// dropped and re-added between shard builds) and concatenate them as one. +/// +/// Absent when the storage has no covering columns. +pub const COVERING_FIELD_IDS_KEY: &str = "covering_field_ids"; + #[derive(Debug)] pub struct QueryScratch { pub distances: Vec, @@ -420,6 +433,17 @@ pub trait VectorStore: Send + Sync + Sized + Clone { where Self: 'a; + /// This storage's own column names: the row id, its quantization code + /// columns, and any legacy bookkeeping column it carries. Everything else in + /// the schema is a covering ("included") column. + /// + /// Required rather than defaulted on purpose: only the storage knows which of + /// its columns are code columns, and a wrong answer here silently mislabels + /// covering data (e.g. it would mistake RaBitQ's code columns for covering + /// ones). Making it a required associated const means a new storage cannot + /// forget to answer. + const INTERNAL_COLUMNS: &'static [&'static str]; + fn as_any(&self) -> &dyn Any; fn schema(&self) -> &SchemaRef; @@ -447,14 +471,9 @@ pub trait VectorStore: Send + Sync + Sized + Clone { /// Field indices of the "included"/covering columns: the extra columns /// stored alongside the row id and quantization code so a covered query can - /// skip the take from the base table. - /// - /// The default is empty (no covering). A storage that supports covering must - /// override this, since only it knows which of its columns are quantization - /// code columns versus included columns — inferring that generically by name - /// is unsafe (e.g. it would mistake RaBitQ's code columns for covering ones). + /// skip the take from the base table. Derived from [`Self::INTERNAL_COLUMNS`]. fn covering_field_indices(&self) -> Vec { - Vec::new() + covering_field_indices_excluding(self.schema().as_ref(), Self::INTERNAL_COLUMNS) } /// `[_rowid, ]` for the whole storage. Captured while a diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index ada522328bb..616432b6905 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -426,6 +426,10 @@ pub(crate) async fn build_index_metadata_from_segments( covered_fragments |= segment.fragment_bitmap().clone(); } + // The declaration is a dependency list, so validate it against the dataset schema and + // prune stale fragment coverage before committing. Physical covering payload is not a + // commit precondition: readers independently verify every segment's storage capability + // and fall back to a base-table take for any declared field that is absent. prune_stale_segment_coverage(dataset, &mut segments, false, false).await?; let new_indices = futures::stream::iter(segments.into_iter().map(|segment| async move { @@ -1258,6 +1262,24 @@ pub(crate) async fn remap_index( let created_index = match generic.index_type() { it if it.is_scalar() => { + // Withdraw rather than remap a covered scalar index. The vector storages' + // `remap` carries their covering columns through row-aligned, which is what + // lets a covered vector index be remapped in place; `BTreeIndex::remap` and + // friends rewrite only `(value, row_id)`. Remapping would republish a + // declaration the new files do not back -- and because carried fields live in + // `fields`, the index would then shed fragment coverage on every update to a + // column it never held. Withdrawing follows the same shape as the + // `can_remap()` check below and the open failure above: a table-level + // compaction degrades one index rather than failing outright. + if !matched.covering_fields.is_empty() { + log::warn!( + "Withdrawing covered index '{}' during remap: its storage format does \ + not carry covering columns through a remap, so the declaration would \ + outlive the payload.", + matched.name + ); + return Ok(RemapResult::Drop); + } let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_id)?; let scalar_index = dataset @@ -8973,11 +8995,12 @@ mod tests { ); } - /// A covering segment's `covering_fields` must survive `commit_existing_index_segments`. - /// The segment files still hold the payload, so dropping the declaration would silently - /// fall covered queries back to a base-table take. + /// The logical covering declaration and the physical payload are independent at commit. + /// Transitional segments may declare a dependency before their writer emits the values, + /// while physical payload without a declaration remains an unused implementation detail. #[tokio::test] - async fn test_commit_existing_index_segments_preserves_covering_fields() { + async fn test_commit_existing_index_segments_preserves_covering_declaration_independently() { + use crate::index::vector::VectorIndexParams; use lance_datagen::{BatchCount, RowCount, array}; let test_dir = tempfile::tempdir().unwrap(); @@ -8989,58 +9012,84 @@ mod tests { "vector", array::rand_vec::(8.into()), ) - .into_reader_rows(RowCount::from(20), BatchCount::from(2)); - - let mut dataset = Dataset::write( - reader, - test_uri, - Some(WriteParams { - max_rows_per_file: 10, - max_rows_per_group: 10, - ..Default::default() - }), - ) - .await - .unwrap(); - + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); let id_field_id = dataset.schema().field("id").unwrap().id; - // A real covered segment: its auxiliary storage physically holds the "id" - // payload, which the commit-boundary cross-check verifies. - let mut params = crate::index::vector::VectorIndexParams::ivf_flat( - 1, - lance_linalg::distance::MetricType::L2, - ); - params.covering_columns(vec!["id".to_string()]); - let seg = dataset + + // Direction 1: storage built WITHOUT covering, declaration fabricated. + let params = VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + let mut plain_segment = dataset .create_index_builder(&["vector"], IndexType::Vector, ¶ms) .name("vector_idx".to_string()) .execute_uncommitted() .await .unwrap(); - assert_eq!(seg.covering_fields, vec![id_field_id]); - + assert!(plain_segment.covering_fields.is_empty()); + // Carried columns live at the end of `fields`, so fabricating a declaration means + // extending both lists, not just `covering_fields`. + plain_segment.fields.push(id_field_id); + plain_segment.covering_fields = vec![id_field_id]; dataset .commit_existing_index_segments( "vector_idx", "vector", - vec![segment_from_metadata(&seg)], + vec![segment_from_metadata(&plain_segment)], ) .await .unwrap(); + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!(committed[0].covering_fields, vec![id_field_id]); + // Direction 2: storage built WITH covering, declaration dropped. + let mut covered_params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + covered_params.covering_columns(vec!["id".to_string()]); + let mut covered_segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, &covered_params) + .name("vector_idx".to_string()) + .replace(true) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(covered_segment.covering_fields, vec![id_field_id]); + // Dropping the declaration drops the carried suffix from `fields` too, leaving a + // segment that claims to key on `vector` alone while its storage still holds `id`. + covered_segment.fields.pop(); + covered_segment.covering_fields = Vec::new(); + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&covered_segment)], + ) + .await + .unwrap(); + let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert!(committed[0].covering_fields.is_empty()); + + // Sanity: the honest covered segment commits fine, and the declaration survives the + // round trip. That matters beyond bookkeeping -- the segment files still hold the + // payload, so losing the declaration would silently drop covered queries back to a + // base-table take. + covered_segment.fields.push(id_field_id); + covered_segment.covering_fields = vec![id_field_id]; + dataset + .commit_existing_index_segments( + "vector_idx", + "vector", + vec![segment_from_metadata(&covered_segment)], + ) + .await + .unwrap(); let committed = dataset.load_indices_by_name("vector_idx").await.unwrap(); assert_eq!(committed.len(), 1); - assert_eq!( - committed[0].covering_fields, - vec![id_field_id], - "committed segment must preserve its covering `covering_fields`" - ); + assert_eq!(committed[0].covering_fields, vec![id_field_id]); } /// A second commit whose incoming segment covers a different column set than an /// existing, disjoint segment we would *retain* must be rejected: all segments of one - /// logical index must agree on their covering, or the committed `covering_fields` - /// would misdescribe some segment's storage. + /// logical index must agree on their declaration so they expose one dependency set and + /// one declaration-ordered query schema. Their physical capabilities may still differ. #[tokio::test] async fn test_commit_existing_index_segments_rejects_inconsistent_retained_covering() { use lance_datagen::{BatchCount, RowCount, array}; @@ -9421,6 +9470,7 @@ mod tests { /// `commit_existing_index_segments`. #[tokio::test] async fn test_build_index_metadata_from_segments_accepts_carried_fields() { + use crate::index::vector::VectorIndexParams; use lance_datagen::{BatchCount, RowCount, array}; let test_dir = tempfile::tempdir().unwrap(); @@ -9430,25 +9480,28 @@ mod tests { "vector", array::rand_vec::(8.into()), ) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - let dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .into_reader_rows(RowCount::from(256), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) .await .unwrap(); let vector_field_id = dataset.schema().field("vector").unwrap().id; let id_field_id = dataset.schema().field("id").unwrap().id; - let mut metadata = write_vector_segment_metadata( - &dataset, - "vector_idx", - vector_field_id, - Uuid::new_v4(), - [0_u32], - b"segment", - ) - .await; - // Carry `id` alongside the keyed `vector` field. - metadata.fields = vec![vector_field_id, id_field_id]; - metadata.covering_fields = vec![id_field_id]; + // A genuinely covered segment, not a hand-edited declaration: the commit boundary + // now cross-checks `covering_fields` against the payload in the segment's auxiliary + // storage, so `id` has to really be there. + let mut params = + VectorIndexParams::ivf_pq(2, 8, 4, lance_linalg::distance::MetricType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + let metadata = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vector_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + // Carried `id` sits at the end of `fields`, behind the keyed `vector` field. + assert_eq!(metadata.fields, vec![vector_field_id, id_field_id]); + assert_eq!(metadata.covering_fields, vec![id_field_id]); let new_indices = build_index_metadata_from_segments( &dataset, @@ -9710,6 +9763,109 @@ mod tests { ); } + /// `covering_fields` is a format-level declaration shared by every index type. A + /// multi-field `fields` with a carried suffix must therefore remain legal even when + /// that index implementation has no physical covering capability yet. + /// A covered SCALAR index must be withdrawn by remap, not remapped. + /// + /// `remap_index` stopped withdrawing covered indexes on the grounds that "each storage + /// format's `remap` preserves its covering columns row-aligned". That holds for the four + /// vector storages; `BTreeIndex::remap` rewrites only `(value, row_id)`. Remapping one + /// republishes a declaration nothing backs, and because carried fields live in `fields` + /// the index then loses fragment coverage on every update to a column it never held. + #[tokio::test] + async fn test_remap_withdraws_a_covered_scalar_index() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("other", array::step::()) + .into_reader_rows(RowCount::from(64), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let other_field_id = dataset.schema().field("other").unwrap().id; + + let btree_params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + dataset + .create_index( + &["id"], + IndexType::BTree, + Some("id_idx".into()), + &btree_params, + true, + ) + .await + .unwrap(); + + // Actually COMMIT the covered declaration -- mutating a loaded copy proves nothing, + // because `remap_index` re-loads the indices from the dataset itself. + crate::utils::test::covering::declare_covering(&mut dataset, "id", "other").await; + let indices = dataset.load_indices().await.unwrap(); + let committed = indices.iter().find(|i| i.name == "id_idx").unwrap(); + assert_eq!( + committed.covering_fields, + vec![other_field_id], + "precondition: the covered declaration must be committed, or this test is vacuous" + ); + assert_eq!(committed.fields, vec![id_field_id, other_field_id]); + let index_id = committed.uuid; + + let result = remap_index(&dataset, &index_id, &RowAddrRemap::empty()).await; + assert!( + matches!(result, Ok(RemapResult::Drop)), + "a covered scalar index must be withdrawn rather than remapped without its \ + payload; got {result:?}" + ); + } + + #[tokio::test] + async fn test_build_index_metadata_from_segments_accepts_carried_fields_on_scalar() { + use lance_datagen::{BatchCount, RowCount, array}; + + let test_dir = tempfile::tempdir().unwrap(); + let reader = lance_datagen::gen_batch() + .col("id", array::step::()) + .col("other", array::step::()) + .into_reader_rows(RowCount::from(64), BatchCount::from(1)); + let mut dataset = Dataset::write(reader, test_dir.path().to_str().unwrap(), None) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let other_field_id = dataset.schema().field("other").unwrap().id; + + let btree_params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + let mut metadata = dataset + .create_index_builder(&["id"], IndexType::BTree, &btree_params) + .name("id_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + // No producer builds covered scalar indexes (`covering_columns` lives only on + // `VectorIndexParams`), so the declaration is hand-constructed on a real BTree + // segment. That is the shape a distributed caller can hand to the public + // `commit_existing_index_segments`, which is what this guard has to tolerate. + metadata.fields = vec![id_field_id, other_field_id]; + metadata.covering_fields = vec![other_field_id]; + + let new_indices = build_index_metadata_from_segments( + &dataset, + "id_idx", + id_field_id, + vec![segment_from_metadata(&metadata)], + ) + .await + .expect("a carried-field declaration must not require physical payload at commit"); + + assert_eq!(new_indices.len(), 1); + assert_eq!(new_indices[0].fields, vec![id_field_id, other_field_id]); + assert_eq!(new_indices[0].covering_fields, vec![other_field_id]); + } + /// A segment keyed on the wrong field must still be rejected. Calls /// `build_index_metadata_from_segments` directly so the assertion is /// evidence for *this* guard specifically -- going through diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index ce37c91fe54..dffb525927d 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -72,6 +72,29 @@ use crate::{Error, Result, dataset::Dataset, index::pb::vector_index_stage::Stag pub const LANCE_VECTOR_INDEX: &str = "__lance_vector_index"; +/// Names a user column can never take inside vector index storage: the union of every +/// storage's exported internal set (row id, partition id, each quantizer's code/factor +/// columns -- extending any `*_INTERNAL_COLUMNS` const extends this set automatically) +/// plus the pipeline-transient names that never reach storage (the distance column and +/// the IVF partition transform's `__centroid_dist`). Used both to reject reserved names +/// in `covering_columns` at create time and to classify which columns of an index +/// storage schema are covering ("included") payload. +pub(crate) static RESERVED_STORAGE_COLUMNS: std::sync::LazyLock< + std::collections::HashSet<&'static str>, +> = std::sync::LazyLock::new(|| { + lance_index::vector::pq::storage::PQ_INTERNAL_COLUMNS + .iter() + .chain(lance_index::vector::sq::storage::SQ_INTERNAL_COLUMNS) + .chain(lance_index::vector::flat::storage::FLAT_INTERNAL_COLUMNS) + .chain(lance_index::vector::bq::storage::RABIT_INTERNAL_COLUMNS) + .chain(&[ + lance_index::vector::DIST_COL, + lance_index::vector::CENTROID_DIST_COLUMN, + ]) + .copied() + .collect() +}); + /// A materialized snapshot of one logical vector index and all of its segments. #[derive(Debug)] pub struct LogicalVectorIndex { @@ -619,6 +642,79 @@ async fn prepare_vector_segment_build( Ok((element_type, index_type, ivf_params, shuffler)) } +/// Validate `covering_columns` (covering columns) at the API boundary, before any build. +/// Rejects non-V3 format, nested/dotted names, the indexed vector column itself, names that +/// collide with a build pipeline's internal storage/transform columns, duplicates, missing +/// columns, and blob columns. Shared by the single-node and distributed vector build paths so +/// the reserved-name list lives in one place. No-op when no covering columns are configured. +fn validate_covering_columns( + params: &VectorIndexParams, + dataset: &Dataset, + column: &str, +) -> Result<()> { + if params.covering_columns.is_empty() { + return Ok(()); + } + + // Covering ("included") columns require the V3 index file format: only the V3 + // covering-aware storages preserve the extra columns row-aligned and report them via + // `covering_field_indices`. Every vector index type supports covering at V3, so reject + // only legacy/non-V3 builds up front -- otherwise the build silently ignores the option + // while `covering_fields` still lands in the index metadata, and the search exec then + // declares a covered schema the storage cannot emit. + if !matches!(params.version, IndexFileVersion::V3) { + return Err(Error::invalid_input(format!( + "covering_columns (covering columns) requires index file version V3, but got {:?}", + params.version + ))); + } + + // A covering column is stored inline in the index, row-aligned with the code, and + // projected by name from the source batch. Reject columns that cannot satisfy that + // contract up front (otherwise the build fails deep in projection, or -- worse -- + // stores unusable data): nested/dotted paths (the covering gather projects only + // top-level names) and blob columns (which store out-of-line descriptors, not inline + // data). + let mut seen_include = std::collections::HashSet::with_capacity(params.covering_columns.len()); + for name in ¶ms.covering_columns { + if name.contains('.') { + return Err(Error::invalid_input(format!( + "covering_columns: nested/dotted covering column '{name}' is not supported; only \ + top-level columns can be covered" + ))); + } + if name == column { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is the indexed vector column itself; it \ + is stored as the quantization code, not a covered payload" + ))); + } + if RESERVED_STORAGE_COLUMNS.contains(&name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' collides with a reserved index storage \ + column name" + ))); + } + if !seen_include.insert(name.as_str()) { + return Err(Error::invalid_input(format!( + "covering_columns: duplicate covering column '{name}'" + ))); + } + let field = dataset.schema().field(name).ok_or_else(|| { + Error::invalid_input(format!( + "covering_columns: covering column '{name}' does not exist in the dataset schema" + )) + })?; + if field.is_blob() { + return Err(Error::invalid_input(format!( + "covering_columns: covering column '{name}' is a blob column; blob columns cannot be \ + covered by a vector index" + ))); + } + } + Ok(()) +} + /// Build a Distributed Vector Index for specific fragments #[allow(clippy::too_many_arguments)] #[instrument(level = "debug", skip(dataset))] @@ -632,18 +728,15 @@ pub(crate) async fn build_distributed_vector_index( fragment_ids: &[u32], progress: Arc, ) -> Result<(Uuid, Vec)> { - // The distributed shard builders and the distributed merger have no covering-column - // plumbing: each shard would write storage without the payload while the committed - // metadata still advertises `covering_fields`, desyncing the two. Reject covering on - // this path up front rather than publishing an inconsistent index. - if !params.covering_columns.is_empty() { - return Err(Error::invalid_input( - "covering_columns (covering columns) are not supported for distributed vector index \ - builds (precomputed IVF centroids with a fragment subset). Build the covered index \ - without a fragment restriction / precomputed IVF." - .to_string(), - )); - } + // Each shard reuses `IvfIndexBuilder`, which threads `covering_columns` through its + // projection and per-partition storage exactly like the single-node build (the shard is + // just a precomputed-IVF, fragment-filtered build). Validate the covering columns up + // front, same as the single-node path, then thread them into each shard builder below so + // the segment storage carries the payload the committed `covering_fields` advertises. + // The cross-shard segment merger carries covering too (`merge_partial_vector_auxiliary_files` + // appends the covering fields to each writer's schema and rejects shards whose covering sets + // disagree), so both per-segment commit and merged builds are covered. + validate_covering_columns(params, dataset, column)?; let (element_type, index_type, ivf_params, shuffler) = prepare_vector_segment_build( dataset, column, @@ -721,6 +814,7 @@ pub(crate) async fn build_distributed_vector_index( frag_reuse_index, )? .with_ivf(ivf_model) + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -742,6 +836,7 @@ pub(crate) async fn build_distributed_vector_index( frag_reuse_index, )? .with_ivf(ivf_model) + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -791,6 +886,7 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep PQ codes in row-major layout. // A single transpose is performed in the distributed merge stage. .with_transpose(false) + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -818,6 +914,7 @@ pub(crate) async fn build_distributed_vector_index( (), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -846,6 +943,7 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -864,6 +962,7 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -906,6 +1005,7 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep PQ codes in row-major layout. // A single transpose is performed in the distributed merge stage. .with_transpose(false) + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -937,6 +1037,7 @@ pub(crate) async fn build_distributed_vector_index( hnsw_params.clone(), frag_reuse_index, )? + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -969,6 +1070,7 @@ pub(crate) async fn build_distributed_vector_index( // For distributed shards, keep RQ codes in row-major layout. // A single packing pass is performed in the distributed merge stage. .with_transpose(false) + .with_covering_columns(params.covering_columns.clone()) .with_fragment_filter(fragment_filter) .with_progress(progress.clone()) .build() @@ -1057,86 +1159,7 @@ async fn build_vector_index_impl( .await?; let stages = ¶ms.stages; - // Covering ("included") columns require the V3 index file format: only the V3 - // covering-aware storages preserve the extra columns row-aligned and report them via - // `covering_field_indices`. Every vector index type supports covering at V3, so reject - // only legacy/non-V3 builds up front -- otherwise the build silently ignores the option - // while `covering_fields` still lands in the index metadata, and the search exec then - // declares a covered schema the storage cannot emit. - if !params.covering_columns.is_empty() && !matches!(params.version, IndexFileVersion::V3) { - return Err(Error::invalid_input(format!( - "covering_columns (covering columns) requires index file version V3, but got {:?}", - params.version - ))); - } - - // A covering column is stored inline in the index, row-aligned with the code, and - // projected by name from the source batch. Reject columns that cannot satisfy that - // contract up front (otherwise the build fails deep in projection, or -- worse -- - // stores unusable data): nested/dotted paths (the covering gather projects only - // top-level names) and blob columns (which store out-of-line descriptors, not inline - // data). - // Names that collide with a build pipeline's own internal columns. Covering one would - // advertise a column in `covering_fields` that the storage's `covering_field_indices` - // excludes (so a covered query declares a column storage never emits), or fail deep in - // the quantizer transform. This check runs before the `match index_type` below, so it - // must list the *union* of every pipeline's internal names -- the row id, distance, - // partition id, each quantizer's code column, RaBitQ's extended-code and per-row factor - // columns, and the IVF partition transform's transient `__centroid_dist`. - const RESERVED_STORAGE_COLUMNS: &[&str] = &[ - lance_core::ROW_ID, - lance_index::vector::DIST_COL, - lance_index::vector::PART_ID_COLUMN, - lance_index::vector::CENTROID_DIST_COLUMN, - lance_index::vector::PQ_CODE_COLUMN, - lance_index::vector::SQ_CODE_COLUMN, - lance_index::vector::flat::storage::FLAT_COLUMN, - lance_index::vector::bq::storage::RABIT_CODE_COLUMN, - lance_index::vector::bq::storage::RABIT_EX_CODE_COLUMN, - lance_index::vector::bq::storage::RABIT_BLOCKED_EX_CODE_COLUMN, - lance_index::vector::bq::transform::ADD_FACTORS_COLUMN, - lance_index::vector::bq::transform::SCALE_FACTORS_COLUMN, - lance_index::vector::bq::transform::ERROR_FACTORS_COLUMN, - lance_index::vector::bq::transform::EX_ADD_FACTORS_COLUMN, - lance_index::vector::bq::transform::EX_SCALE_FACTORS_COLUMN, - ]; - let mut seen_include = std::collections::HashSet::with_capacity(params.covering_columns.len()); - for name in ¶ms.covering_columns { - if name.contains('.') { - return Err(Error::invalid_input(format!( - "covering_columns: nested/dotted covering column '{name}' is not supported; only \ - top-level columns can be covered" - ))); - } - if name == column { - return Err(Error::invalid_input(format!( - "covering_columns: covering column '{name}' is the indexed vector column itself; it \ - is stored as the quantization code, not a covered payload" - ))); - } - if RESERVED_STORAGE_COLUMNS.contains(&name.as_str()) { - return Err(Error::invalid_input(format!( - "covering_columns: covering column '{name}' collides with a reserved index storage \ - column name" - ))); - } - if !seen_include.insert(name.as_str()) { - return Err(Error::invalid_input(format!( - "covering_columns: duplicate covering column '{name}'" - ))); - } - let field = dataset.schema().field(name).ok_or_else(|| { - Error::invalid_input(format!( - "covering_columns: covering column '{name}' does not exist in the dataset schema" - )) - })?; - if field.is_blob() { - return Err(Error::invalid_input(format!( - "covering_columns: covering column '{name}' is a blob column; blob columns cannot be \ - covered by a vector index" - ))); - } - } + validate_covering_columns(params, dataset, column)?; match index_type { IndexType::IvfFlat => match element_type { @@ -2075,6 +2098,11 @@ pub async fn initialize_vector_index( }) .collect::>>()?; params.covering_columns(covering_columns); + // Validate against the TARGET dataset, not the source: the names resolve there, but + // the column they resolve to may be a different kind (a blob, a reserved storage name, + // a non-V3 build). Without this, the two ordinary build entry points would reject a + // covering set that this cross-dataset path silently accepts. + validate_covering_columns(¶ms, target_dataset, column_name)?; let new_uuid = Uuid::new_v4(); let frag_reuse_index = target_dataset @@ -2859,6 +2887,97 @@ mod tests { } } + /// `initialize_vector_index` re-resolves the source index's covering columns + /// against the TARGET schema, so a name that was valid to cover in the source can + /// resolve to something that is not coverable in the target. It must run the same + /// validation the two ordinary build entry points do, or it silently materializes + /// (here) blob descriptor structs as the covered payload. + #[tokio::test] + async fn test_initialize_vector_index_validates_covering_against_target() { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + // Source: `payload` is an ordinary binary column, perfectly coverable. + let source_reader = lance_datagen::gen_batch() + .col("payload", array::rand_type(&ArrowDataType::LargeBinary)) + .col("vector", array::rand_vec::(32.into())) + .into_reader_rows(RowCount::from(300), BatchCount::from(1)); + let mut source_dataset = Dataset::write(source_reader, &source_uri, None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 20); + params.covering_columns(vec!["payload".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + + // Target: same column name, same arrow type -- but declared a blob, so it + // stores out-of-line descriptors rather than inline data. + let rows = 300usize; + let payload: arrow_array::ArrayRef = + Arc::new(arrow_array::LargeBinaryArray::from_iter_values( + (0..rows).map(|i| format!("v{i}").into_bytes()), + )); + let vectors: arrow_array::ArrayRef = Arc::new( + arrow_array::FixedSizeListArray::try_new_from_values( + arrow_array::Float32Array::from(vec![0.5f32; rows * 32]), + 32, + ) + .unwrap(), + ); + let blob_field = Field::new("payload", ArrowDataType::LargeBinary, true).with_metadata( + [(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())] + .into_iter() + .collect(), + ); + let target_schema = Arc::new(ArrowSchema::new(vec![ + blob_field, + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(target_schema.clone(), vec![payload, vectors]).unwrap(); + let target_reader = + arrow_array::RecordBatchIterator::new(vec![Ok(batch)], target_schema.clone()); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + assert!( + target_dataset.schema().field("payload").unwrap().is_blob(), + "test setup must give the target a blob 'payload', or this test proves nothing" + ); + + let err = initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .expect_err("covering a blob column in the target must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("blob") && msg.contains("payload"), + "error should name the offending blob covering column; was: {msg}" + ); + } + /// Copying/importing a covered index into another dataset (initialize) must /// rebuild the target storage WITH the covering columns, not just copy the /// `covering_fields` metadata -- otherwise the target advertises covering @@ -3253,47 +3372,6 @@ mod tests { ); } - /// The distributed build + merge path has no covering-column plumbing, so it must - /// reject `covering_columns` up front rather than publish an index whose metadata - /// advertises columns the shard files do not contain. - #[tokio::test] - async fn test_build_distributed_rejects_covering_columns() { - let test_dir = TempStrDir::default(); - let uri = format!("{}/ds", test_dir.as_str()); - - let reader = lance_datagen::gen_batch() - .col("id", array::step::()) - .col("vector", array::rand_vec::(32.into())) - .into_reader_rows(RowCount::from(128), BatchCount::from(1)); - let dataset = Dataset::write(reader, &uri, None).await.unwrap(); - - let mut params = VectorIndexParams::ivf_flat(4, MetricType::L2); - params.covering_columns(vec!["id".to_string()]); - - let result = build_distributed_vector_index( - &dataset, - "vector", - "vector_dist", - Uuid::new_v4(), - ¶ms, - None, - &[0], - noop_progress(), - ) - .await; - - assert!( - matches!(&result, Err(Error::InvalidInput { .. })), - "distributed build must reject covering_columns, got {:?}", - result - ); - let msg = result.unwrap_err().to_string(); - assert!( - msg.contains("covering_columns") && msg.contains("distributed"), - "error should mention covering_columns and the distributed path; got: {msg}" - ); - } - #[tokio::test] async fn test_build_distributed_empty_fragment_ids() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 62db4a77e31..f90d9b89b2b 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -43,7 +43,7 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::quantizer::{QuantizerMetadata, QuantizerStorage}; use lance_index::vector::shared::{SupportedIvfIndexType, write_unified_ivf_and_index_metadata}; -use lance_index::vector::storage::STORAGE_METADATA_KEY; +use lance_index::vector::storage::{COVERING_FIELD_IDS_KEY, STORAGE_METADATA_KEY}; use lance_index::vector::transform::Flatten; use lance_index::vector::v3::shuffler::{ DEFAULT_PARTITION_WINDOW_BYTES, EmptyReader, IvfShufflerReader, create_ivf_shuffler, @@ -86,6 +86,7 @@ use crate::Dataset; use crate::dataset::ProjectionRequest; use crate::dataset::index::dataset_format_version; use crate::index::append::build_old_data_filter; +use crate::index::vector::RESERVED_STORAGE_COLUMNS; use crate::index::vector::bounded_partition_stream::{ BoundedPartitionStream, Budgeted, OrderedPartitionResults, WeightedJob, }; @@ -138,6 +139,40 @@ fn append_covering_by_row_id( Ok(batch) } +/// Drop covering ("included") columns that the index metadata does not declare. +/// +/// A pre-covering writer can clear `covering_fields` while the auxiliary storage keeps the +/// payload: every commit re-serializes the whole index list, and that writer's +/// `pb::IndexMetadata` has no field 11, so prost drops the declaration. +/// `FLAG_COVERED_INDEX_METADATA` fences pre-covering builds off the dataset, but the +/// fence is best-effort -- released clients consult writer flags only on the +/// append/overwrite path, so delete/update/raw commits can still slip through it. Existing +/// partition storage then still carries those columns while freshly scanned rows do not +/// (the scan projects only what the metadata declares), and `StorageBuilder::build` rejects +/// the width mismatch, leaving a readable index that can never be merge-optimized again. +/// The manifest declaration is the source of truth here, exactly as it is on the read path, +/// so narrow storage down to it rather than widening the fresh side -- widening would +/// re-materialize a payload the metadata does not advertise. +/// +/// A no-op in the healthy case: storage and `declared` agree, so nothing is projected away. +fn drop_undeclared_covering(batch: RecordBatch, declared: &[String]) -> Result { + let schema = batch.schema(); + let keep: Vec = schema + .fields() + .iter() + .enumerate() + .filter(|(_, field)| { + RESERVED_STORAGE_COLUMNS.contains(field.name().as_str()) + || declared.iter().any(|name| name == field.name()) + }) + .map(|(idx, _)| idx) + .collect(); + if keep.len() == batch.num_columns() { + return Ok(batch); + } + Ok(batch.project(&keep)?) +} + /// Build a new centroid array that incorporates the results of partition splits. /// /// For each `(part_idx, centroid1, centroid2)` in `splits`: @@ -796,6 +831,7 @@ impl IvfIndexBuilder return Err(Error::invalid_input("dataset not set before shuffling")); }; + let mut from_precomputed_buffers = false; let stream = match self .ivf_params .as_ref() @@ -812,6 +848,7 @@ impl IvfIndexBuilder .to_string(), )); } + from_precomputed_buffers = true; let uri = to_local_path(uri); // the uri points to data directory, // so need to trim the "data" suffix for reading the dataset @@ -853,10 +890,16 @@ impl IvfIndexBuilder } }; - if let Some((row_id_idx, _)) = stream.schema().column_with_name("row_id") { - // When using precomputed shuffle buffers we can't use the column name _rowid - // since it is reserved. So we tolerate `row_id` as well here (and rename it - // to _rowid to match the non-precomputed path) + // Scoped to the precomputed-buffer path on purpose. Those buffers cannot name a + // column `_rowid` (reserved), so they materialize it as `row_id` and it is renamed + // back here to match the non-precomputed path. The ordinary scan above already + // produces a real `_rowid` via `with_row_id()`, and it now also projects the + // covering columns -- so renaming there would rewrite a *user* column named + // `row_id` on top of the real one and fail index creation with + // `Duplicate field name "_rowid" in schema`. + if from_precomputed_buffers + && let Some((row_id_idx, _)) = stream.schema().column_with_name("row_id") + { self.shuffle_data(Some(Self::rename_row_id(stream, row_id_idx))) .await?; } else { @@ -1146,6 +1189,10 @@ impl IvfIndexBuilder // Covering builds may re-scan a pruned fragment that still lives in old // storage; drop the stale duplicate row ids during the partition read. let dedup_existing = !self.covering_columns.is_empty(); + // The covering set this build declares. Existing storage may be WIDER (see + // `drop_undeclared_covering`); narrow it so every batch handed to + // `StorageBuilder::build` agrees on width. + let declared_covering = Arc::new(self.covering_columns.clone()); let partition_adjustment = Arc::new(partition_adjustment); let build_iter = assign_batches @@ -1161,6 +1208,7 @@ impl IvfIndexBuilder let column = column.clone(); let frag_reuse_index = frag_reuse_index.clone(); let dedup_existing = dedup_existing; + let declared_covering = declared_covering.clone(); let partition_adjustment = partition_adjustment.clone(); async move { let (is_affected, split_reader) = match partition_adjustment.as_ref() { @@ -1218,6 +1266,14 @@ impl IvfIndexBuilder loss += extra_loss; } + // Existing storage can carry covering columns this build does not + // declare; narrow before the fresh rows join them below. + + batches = batches + .into_iter() + .map(|batch| drop_undeclared_covering(batch, &declared_covering)) + .collect::>>()?; + spawn_cpu(move || { // Apply assign_batch for join operations (splits no // longer use assign_batches) @@ -1685,6 +1741,39 @@ impl IvfIndexBuilder .collect() } + /// The *source dataset* field ids of the covering columns, in declaration order. + /// Stamped into the storage file (see [`COVERING_FIELD_IDS_KEY`]) so a distributed + /// merge can tell shards that cover the same logical fields from shards whose + /// covering columns merely share a name and type. + /// + /// Fails exactly where [`Self::covering_arrow_fields`] fails, and for the same reason: + /// the two must agree in arity or the storage file gets N covering columns and a stamp + /// of a different length, which makes `physical_covering_fields_from_schema` withdraw + /// covering for that index permanently and silently. + fn covering_field_ids(&self) -> Result> { + if self.covering_columns.is_empty() { + return Ok(Vec::new()); + } + let Some(ds) = self.dataset.as_ref() else { + return Err(Error::invalid_input( + "dataset not set before resolving covering field ids".to_string(), + )); + }; + self.covering_columns + .iter() + .map(|name| { + ds.schema() + .field(name) + .map(|field| field.id) + .ok_or_else(|| { + Error::invalid_input(format!( + "include column '{name}' not found in dataset schema" + )) + }) + }) + .collect() + } + #[instrument(name = "merge_partitions", level = "debug", skip_all)] async fn merge_partitions( &mut self, @@ -1722,13 +1811,21 @@ impl IvfIndexBuilder // before `_rowid` and make every downstream reader that trusts column position // (e.g. the default `QuantizerStorage::remap`) misread row ids. let covering_fields = self.covering_arrow_fields()?; + // Every storage batch is reordered into this declared schema before it is written: + // `FileWriter` picks columns by name when encoding but checks nullability by + // position, so a batch ordered `[covering…, _rowid, code]` would be checked against + // the wrong fields and a null-bearing covering column rejected as if it were the + // (non-nullable) code column. + let mut storage_arrow_schema: Option = None; let mut storage_writer = if is_flat { None } else { let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; fields.extend(quantizer.extra_fields()); fields.extend(covering_fields.iter().cloned()); - let storage_schema: Schema = (&arrow_schema::Schema::new(fields)).try_into()?; + let arrow_schema = Arc::new(arrow_schema::Schema::new(fields)); + let storage_schema: Schema = arrow_schema.as_ref().try_into()?; + storage_arrow_schema = Some(arrow_schema); Some(file_versions::create_writer( self.format_version, self.store.create(&storage_path).await?, @@ -1834,8 +1931,9 @@ impl IvfIndexBuilder .clone(); let mut fields = vec![ROW_ID_FIELD.as_ref().clone(), flat_field]; fields.extend(covering_fields.iter().cloned()); - let storage_schema: Schema = - (&arrow_schema::Schema::new(fields)).try_into()?; + let arrow_schema = Arc::new(arrow_schema::Schema::new(fields)); + let storage_schema: Schema = arrow_schema.as_ref().try_into()?; + storage_arrow_schema = Some(arrow_schema); storage_writer = Some(file_versions::create_writer( self.format_version, self.store.create(&storage_path).await?, @@ -1843,6 +1941,20 @@ impl IvfIndexBuilder writer_options.clone(), )?); } + // Match the declared column order (see `storage_arrow_schema`). + // + // Load-bearing, not tidiness: the writer schema is declared as + // `[_rowid, code, extra..., covering...]` while a covered batch arrives + // as `[covering..., _rowid, code]`, and the file writer encodes columns + // **by name** but checks nullability **by position**. Writing the batch + // unprojected therefore matches a non-nullable declared slot against a + // nullable covering column and kills the build on the first null covering + // value. The pre-covering code derived the writer schema from the batch + // and so could not disagree; declaring it instead -- which is what stops a + // positional row-id read corrupting row ids -- is what made this necessary. + if let Some(schema) = storage_arrow_schema.as_ref() { + batch = batch.project_by_schema(schema)?; + } storage_writer .as_mut() .expect("storage writer must be initialized before write") @@ -1953,6 +2065,17 @@ impl IvfIndexBuilder STORAGE_METADATA_KEY, serde_json::to_string(&storage_partition_metadata)?, ); + let covering_field_ids = self.covering_field_ids()?; + if !covering_field_ids.is_empty() { + storage_writer.add_schema_metadata( + COVERING_FIELD_IDS_KEY, + covering_field_ids + .iter() + .map(|id| id.to_string()) + .collect::>() + .join(","), + ); + } let index_type_str = index_type_string(S::name().try_into()?, Q::quantization_type()); if let Some(idx_type) = SupportedIvfIndexType::from_index_type_str(&index_type_str) { diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 35654acafbe..fd1a8f79174 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -2547,6 +2547,36 @@ mod tests { assert_eq!(capacity.u32, 0); } + /// Vector values laid out as `num_clusters` well-separated clusters: row `r` lands in + /// cluster `r % num_clusters`, centred at `cluster * 1000.0` with a tiny per-row offset + /// so no two rows coincide. + /// + /// **The separation is load-bearing, not cosmetic.** `early_pruning` RAISES + /// `minimum_nprobes` to the number of centroids within `dists[0] * factor` (7.0 for + /// `k` in 2..=10, 81.0 for `k >= 11`). With evenly-spread vectors every centroid + /// survives pruning, so `minimum_nprobes` reaches `maximum_nprobes`, `late_search` + /// returns at its first guard, and the no-rows shortcut -- and therefore the covered + /// recovery path -- is never reached. A test targeting either then passes vacuously + /// against a live bug, which has happened repeatedly on this feature. With these + /// clusters a query at the origin is ~10^6 times closer to cluster 0 than to any + /// other, so pruning yields 1 and unsearched partitions remain. + /// + /// See also `generate_clustered_batch`, the pre-existing generator used by the + /// non-covering partition-split tests; it emits a different schema and is not a + /// drop-in substitute here. + fn clustered_vector_values( + rows: impl IntoIterator, + dim: i32, + num_clusters: usize, + ) -> Vec { + rows.into_iter() + .flat_map(|r| { + let center = (r % num_clusters) as f32 * 1000.0; + (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) + }) + .collect() + } + async fn generate_test_dataset( test_uri: &str, range: Range, @@ -3665,12 +3695,7 @@ mod tests { // Two well-separated clusters so early pruning keeps minimum_nprobes at 1: // the recovery path only fires where the non-covered shortcut would (an // all-partitions-searched query returns only found rows on both). - let values: Vec = (0..n) - .flat_map(|r| { - let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; - (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) - }) - .collect(); + let values: Vec = clustered_vector_values(0..n, dim, 2); let validity: Vec = (0..n).map(|i| i >= n_rare).collect(); let vector = FixedSizeListArray::new( Arc::new(Field::new("item", DataType::Float32, true)), @@ -3787,12 +3812,7 @@ mod tests { let cats: Vec<&str> = (0..n).map(|_| "x").collect(); // Two well-separated clusters so early pruning keeps minimum_nprobes at 1 // (see test_ivf_covered_recovers_null_vector_prefilter_rows). - let values: Vec = (0..n) - .flat_map(|r| { - let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; - (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) - }) - .collect(); + let values: Vec = clustered_vector_values(0..n, dim, 2); let validity: Vec = (0..n).map(|i| i >= n_null).collect(); let vector = FixedSizeListArray::new( Arc::new(Field::new("item", DataType::Float32, true)), @@ -3905,8 +3925,13 @@ mod tests { ])); let make_batch = |ids: Vec, null_rows: &[i32]| { - let n = ids.len(); - let values: Vec = (0..n * dim as usize).map(|i| i as f32 + 1.0).collect(); + // Two well-separated clusters. Recovery is only reachable while + // `minimum_nprobes < maximum_nprobes`, and `early_pruning` raises the minimum to + // the number of centroids within `dists[0] * 7.0` (k = 10). Points strung along + // a line leave those centroids only ~7.5x apart -- a margin thin enough that a + // different kmeans outcome silently disables the path this test exists to cover. + let values: Vec = + clustered_vector_values(ids.iter().map(|id| *id as usize), dim, 2); let validity: Vec = ids.iter().map(|id| !null_rows.contains(id)).collect(); let vector = FixedSizeListArray::new( Arc::new(Field::new("item", DataType::Float32, true)), @@ -4040,13 +4065,8 @@ mod tests { // Two separated clusters so `early_pruning` leaves an unsearched partition and the // late-search shortcut (which arms the covered recovery) is reachable. let make_batch = |ids: Vec| { - let values: Vec = ids - .iter() - .flat_map(|id| { - let center = if id % 2 == 0 { 0.0f32 } else { 1000.0 }; - (0..dim as usize).map(move |d| center + (*id as usize * 4 + d) as f32 * 1e-3) - }) - .collect(); + let values: Vec = + clustered_vector_values(ids.iter().map(|id| *id as usize), dim, 2); RecordBatch::try_new( schema.clone(), vec![ @@ -4156,12 +4176,7 @@ mod tests { // hence the covered recovery) reachable -- with evenly spread vectors early pruning // raises minimum_nprobes to cover every partition and the shortcut never fires. let ids: Vec = (0..n as i32).collect(); - let values: Vec = (0..n) - .flat_map(|r| { - let center = if r % 2 == 0 { 0.0f32 } else { 1000.0 }; - (0..dim as usize).map(move |d| center + (r * dim as usize + d) as f32 * 1e-3) - }) - .collect(); + let values: Vec = clustered_vector_values(0..n, dim, 2); let vector = FixedSizeListArray::new( Arc::new(Field::new("item", DataType::Float32, true)), dim, @@ -4252,6 +4267,195 @@ mod tests { ); } + /// A covered index whose `covering_fields` a pre-covering writer cleared (prost drops + /// unknown proto field 11 on re-serialization; `FLAG_COVERED_INDEX_METADATA` fences + /// this off only best-effort, since released clients consult writer flags solely on + /// the append/overwrite path) -- the degraded state -- must stay MAINTAINABLE, not just + /// readable. Its auxiliary storage still physically carries the payload while freshly + /// scanned rows do not, so combining the two used to fail `StorageBuilder::build`'s width + /// check ("mismatched columns while merging vector storage batches: expected + /// [_rowid, __pq_code, id], got [_rowid, __pq_code]"), leaving an index that could never + /// be merge-optimized again -- recoverable only by dropping and rebuilding it. + /// `OptimizeOptions::append()` never hit this because it does not load existing storage. + #[tokio::test] + async fn test_degraded_covered_index_can_still_be_optimized() { + use crate::dataset::WriteDestination; + use crate::dataset::transaction::Operation; + use arrow_array::{Int32Array, RecordBatchIterator}; + + const DIMS: usize = 16; + const TOTAL: usize = 256; + const NPART: usize = 4; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + + let make = |lo: i32, hi: i32| { + let ids: Vec = (lo..hi).collect(); + let values: Vec = + clustered_vector_values(ids.iter().map(|r| *r as usize), DIMS as i32, NPART); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIMS as i32, + ), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(values), + DIMS as i32, + ) + .unwrap(), + ), + ], + ) + .unwrap(); + (schema, batch) + }; + + let (schema, batch) = make(0, TOTAL as i32); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + test_uri, + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NPART, 8, 4, DistanceType::L2, 2); + params.covering_columns(vec!["id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .unwrap(); + + // Degrade: clear the declaration, leave the payload in the auxiliary file. + let mut cleared = dataset.load_indices_by_name("vector_idx").await.unwrap()[0].clone(); + assert!(!cleared.covering_fields.is_empty()); + cleared.covering_fields = Vec::new(); + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::CreateIndex { + new_indices: vec![cleared], + removed_indices: vec![], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // Fresh rows land in an unindexed fragment; the merge below must combine them with + // the wider existing storage. + let (schema2, batch2) = make(TOTAL as i32, TOTAL as i32 + 128); + dataset + .append(RecordBatchIterator::new([Ok(batch2)], schema2), None) + .await + .unwrap(); + + dataset + .optimize_indices(&OptimizeOptions::merge(1)) + .await + .expect("a metadata-degraded covered index must still merge-optimize"); + + // The rebuilt storage is narrowed to the declaration, so it no longer carries the + // undeclared payload -- and the index still answers queries. + let ctx = load_vector_index_context(&dataset, "vector", "vector_idx").await; + let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + assert!( + storage.batch().column_by_name("id").is_none(), + "merged storage must drop payload the metadata no longer declares" + ); + + let q = Float32Array::from(vec![0.0f32; DIMS]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + assert!(batch.column_by_name("id").is_some()); + } + + /// `row_id` is an ordinary user column name -- `_rowid` is the reserved one -- so it is + /// coverable. The `row_id` -> `_rowid` rename exists only for precomputed shuffle buffers + /// (which cannot name a column `_rowid`), but it used to run on the ordinary scan path + /// too. Once covering started projecting user columns into that scan, covering a column + /// named `row_id` renamed it on top of the real `_rowid` from `with_row_id()` and index + /// creation died with `Duplicate field name "_rowid" in schema`. + #[tokio::test] + async fn test_covering_column_named_row_id_is_supported() { + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, RecordBatchIterator}; + + let dim = 4i32; + let n = 128usize; + + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::Int32, false), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + false, + ), + ])); + let values: Vec = (0..n * dim as usize).map(|i| i as f32).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim) + .unwrap(), + ), + ], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://covering_named_row_id", + None, + ) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_flat(2, DistanceType::L2); + params.covering_columns(vec!["row_id".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, false) + .await + .expect("a covering column named `row_id` must not collide with the virtual _rowid"); + + // And the covered payload is served correctly, not confused with the virtual column. + let q = Float32Array::from(vec![0.0f32; dim as usize]); + let mut scan = dataset.scan(); + scan.nearest("vector", &q, 5).unwrap(); + scan.project(&["row_id"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 5); + let ids = batch["row_id"].as_primitive::(); + let mut got: Vec = ids.values().to_vec(); + got.sort_unstable(); + assert_eq!( + got, + vec![0, 1, 2, 3, 4], + "covered `row_id` payload must be the user's values" + ); + } + /// A covered ("included") struct's payload schema is fixed at index build time. Growing /// that struct via `add_columns` (which commits as `Operation::Merge`) -- even an /// AllNulls, metadata-only child that writes no data file -- would leave the index @@ -4331,6 +4535,14 @@ mod tests { /// Read-side payoff for every vector index type: a query projecting only a covered /// column is satisfied from the index -- no `TakeExec` against the base table -- /// with row-aligned values and sane recall. + /// + /// The fixture offsets `id` away from `_rowid` on purpose. `generate_test_dataset` + /// starts ids at 0, so `id == _rowid` for every row -- and under that fixture the + /// row-alignment assertion below is satisfied by any defect that serves the row id + /// where the covered value belongs, or the reverse. That is exactly the confusion a + /// storage batch ordered `[covering..., _rowid, code]` invites, and it is silent: a + /// `UInt64` covering column substituted for `_rowid` downcasts cleanly. The offset + /// separates the two columns for all seven quantizer families. #[rstest] #[case::pq(VectorIndexParams::ivf_pq(4, 8, 4, DistanceType::L2, 2))] #[case::sq(VectorIndexParams::with_ivf_sq_params( @@ -4373,9 +4585,31 @@ mod tests { #[tokio::test] async fn test_covered_projection_skips_take(#[case] mut params: VectorIndexParams) { const INDEX_NAME: &str = "vector_idx"; + // Disjoint from the row id range (0..NUM_ROWS), so no covered `id` can coincide + // with any row id. + const ID_OFFSET: u64 = 1_000_000; let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); - let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + let (batch, schema) = + generate_batch::(NUM_ROWS, Some(ID_OFFSET), 0.0..1.0, false); + let vectors = Arc::new( + batch + .column_by_name("vector") + .unwrap() + .as_fixed_size_list() + .clone(), + ); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write( + batches, + test_uri, + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); params.covering_columns(vec!["id".to_string()]); dataset @@ -4419,12 +4653,14 @@ mod tests { .expect("row id column") .as_primitive::(); assert_eq!(ids.len(), 10, "should return k=10 rows"); - // Single-fragment, step-id dataset => id == row offset == _rowid, so a correctly - // covered id equals the row id for every returned row (row-aligned, not stale). + // Single-fragment, step-id dataset => id == ID_OFFSET + row offset == ID_OFFSET + + // _rowid, so a correctly covered id is the row id plus the offset for every + // returned row (row-aligned, not stale). The offset also means a row id sourced + // from the covering column instead of `_rowid` cannot satisfy this. for i in 0..ids.len() { assert_eq!( ids.value(i), - row_ids.value(i), + row_ids.value(i) + ID_OFFSET, "covered id must match the row's true id (row {i})" ); } @@ -4702,7 +4938,8 @@ mod tests { // k >> one partition's rows (NUM_ROWS=512 / 4 parts) forces late expansion. scan.nearest("vector", &q, 500).unwrap(); scan.minimum_nprobes(1); - scan.maximum_nprobes(4); + scan.maximum_nprobes(NUM_CLUSTERS); + scan.with_row_id(); scan.project(&["id"]).unwrap(); let plan = scan.explain_plan(true).await.unwrap(); @@ -4711,12 +4948,22 @@ mod tests { "covered projection ['id'] should not require a TakeExec (batch path); plan:\n{plan}" ); let batch = scan.try_into_batch().await.unwrap(); - assert!(batch.column_by_name("id").is_some()); assert!( batch.num_rows() > ROWS_PER_CLUSTER, "late search should have expanded beyond the nearest partition's {ROWS_PER_CLUSTER} rows, got {}", batch.num_rows() ); + // Row alignment, not just presence: ids are assigned in write order over a single + // fragment, so a correctly covered `id` equals its row id on every returned row. + let ids = batch["id"].as_primitive::(); + let row_ids = batch[ROW_ID].as_primitive::(); + for i in 0..ids.len() { + assert_eq!( + ids.value(i), + row_ids.value(i), + "covered payload must stay row-aligned through the batch path" + ); + } } /// A covered query combined with a selective prefilter (scalar index + @@ -5582,6 +5829,63 @@ mod tests { ) .await .unwrap(); + + // ...and the index it produces must be *correct*, not merely buildable. That + // build is where a covered storage batch reaches the fragment-reuse remap in + // its build-time column order -- `[, _rowid, code]`, row id NOT + // at index 0. The search earlier in this test loads the written file instead, + // whose `[_rowid, code, ]` order makes a positional row-id lookup + // right by accident. So without querying this second index, nothing pins + // `remap_row_ids_by_name` locating the row id by name: a positional variant + // remaps the covering column, leaves the real row ids unremapped, and no + // assertion anywhere observes it. + dataset.drop_index(INDEX_NAME).await.unwrap(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + // Pin the ANN path: a flat-KNN fallback would satisfy every assertion below + // without the remapped index being read at all. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the second index; plan was:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the base-table take; plan was:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let ids = batch + .column_by_name("id") + .expect("covered 'id' column must be emitted") + .as_primitive::(); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + let row_id_vec: Vec = row_ids.values().to_vec(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset.take_rows(&row_id_vec, projection).await.unwrap(); + let truth_ids = truth + .column_by_name("id") + .unwrap() + .as_primitive::(); + for i in 0..ids.len() { + assert!( + ids.value(i) >= 100, + "deleted rows (id < 100) must not be returned by the second index" + ); + assert_eq!( + ids.value(i), + truth_ids.value(i), + "row {i} (row_id {}): covered id != true id for the index built \ + while a fragment-reuse index was already present", + row_ids.value(i) + ); + } } /// Fix 2: under query parallelism > 1 the parallel search branch (which uses @@ -6060,6 +6364,60 @@ mod tests { (schema, vec![batch]) } + /// The covering tests' own vector geometry, kept separate from the shared `TWO_FRAG_*` + /// fixture. Global codebook training costs `dim * iters * samples * centroids` and these + /// tests pay it once per case while building and merging several indexes, which is what put + /// them over the one-second local-unit-test budget. What they assert -- that covered values + /// survive every build, merge and lifecycle step and stay row-aligned with the base table -- + /// is independent of quantization resolution, so the vector shrinks instead. `dim / + /// num_sub_vectors` stays 8, as the shared fixture has it, and the codebook stays 8-bit: that + /// is the production default and the width these tests exist to cover. + /// + /// The shared fixture keeps its own dimensions: its tests assert *exact* single-vs-split + /// top-K equality over uniform-random vectors, which a coarser codebook breaks by making + /// distances tie. + const COVERED_DIM: usize = 32; + const COVERED_NUM_SUBVECTORS: usize = 4; + const COVERED_MAX_ITERS: u32 = 4; + + /// Like `make_two_fragment_batches`, but with two covering columns: a non-null `id` and a + /// **nullable** `payload` (every 3rd value null). Covered tests use this so null covering + /// values are exercised through the per-segment build and the cross-shard merge. + fn make_covered_test_batches() -> (Arc, Vec) { + let ids = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + let payload = Arc::new(UInt64Array::from_iter( + (0..TWO_FRAG_NUM_ROWS as u64).map(|v| if v % 3 == 0 { None } else { Some(v + 1000) }), + )); + + // Clustered vectors: uniform-random data is pathological for IVF_PQ recall (the curse of + // dimensionality), which would make a recall gate flaky. Instead pack points into + // well-separated clusters of <= k points each (centers 4.0 apart per dim -> ~23 in L2 at + // `COVERED_DIM`, tiny within-cluster jitter), so a query drawn from a cluster has its whole + // cluster as the unambiguous nearest neighbors and IVF_PQ recall is reliably high. + const CLUSTER_SIZE: usize = 8; + let mut flat = Vec::with_capacity(TWO_FRAG_NUM_ROWS * COVERED_DIM); + for row in 0..TWO_FRAG_NUM_ROWS { + let center = (row / CLUSTER_SIZE) as f32 * 4.0; + let within = (row % CLUSTER_SIZE) as f32; + for d in 0..COVERED_DIM { + flat.push(center + within * 0.002 + d as f32 * 0.00001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), COVERED_DIM as i32) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("payload", DataType::UInt64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, payload, vectors]).unwrap(); + + (schema, vec![batch]) + } + async fn write_dataset_from_batches( test_uri: &str, schema: Arc, @@ -6159,7 +6517,29 @@ mod tests { (ivf_params, pq_params) } - async fn prepare_global_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { + /// `prepare_global_ivf_pq` for the covering fixture: same shared training path, but over + /// `COVERED_DIM` vectors and the covering codebook. + async fn prepare_covered_ivf_pq( + dataset: &Dataset, + vector_column: &str, + ) -> (IvfBuildParams, PQBuildParams) { + prepare_ivf_pq( + dataset, + vector_column, + COVERED_DIM, + TWO_FRAG_NUM_PARTITIONS, + COVERED_NUM_SUBVECTORS, + TWO_FRAG_NUM_BITS, + COVERED_MAX_ITERS, + TWO_FRAG_SAMPLE_RATE, + ) + .await + } + + /// `prepare_global_ivf` for the covering fixture. Distributed builds need every shard to + /// share one set of centroids, so the covered non-PQ cases pre-train them here rather than + /// letting each shard train its own. + async fn prepare_covered_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { let batch = dataset .scan() .project(&[vector_column.to_string()]) @@ -6173,10 +6553,10 @@ mod tests { .as_fixed_size_list(); let dim = vectors.value_length() as usize; - assert_eq!(dim, TWO_FRAG_DIM, "unexpected vector dimension"); + assert_eq!(dim, COVERED_DIM, "unexpected vector dimension"); let values = vectors.values().as_primitive::(); - let kmeans_params = KMeansParams::new(None, TWO_FRAG_MAX_ITERS, 1, DistanceType::L2); + let kmeans_params = KMeansParams::new(None, COVERED_MAX_ITERS, 1, DistanceType::L2); let kmeans = train_kmeans::( values, kmeans_params, @@ -6191,13 +6571,50 @@ mod tests { Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); let mut ivf_params = IvfBuildParams::try_with_centroids(TWO_FRAG_NUM_PARTITIONS, centroids_fsl).unwrap(); - ivf_params.max_iters = TWO_FRAG_MAX_ITERS as usize; + ivf_params.max_iters = COVERED_MAX_ITERS as usize; ivf_params.sample_rate = TWO_FRAG_SAMPLE_RATE; ivf_params } - async fn build_segments_for_fragment_groups( - dataset: &mut Dataset, + async fn prepare_global_ivf(dataset: &Dataset, vector_column: &str) -> IvfBuildParams { + let batch = dataset + .scan() + .project(&[vector_column.to_string()]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let vectors = batch + .column_by_name(vector_column) + .expect("vector column should exist") + .as_fixed_size_list(); + + let dim = vectors.value_length() as usize; + assert_eq!(dim, TWO_FRAG_DIM, "unexpected vector dimension"); + + let values = vectors.values().as_primitive::(); + let kmeans_params = KMeansParams::new(None, TWO_FRAG_MAX_ITERS, 1, DistanceType::L2); + let kmeans = train_kmeans::( + values, + kmeans_params, + dim, + TWO_FRAG_NUM_PARTITIONS, + TWO_FRAG_SAMPLE_RATE, + ) + .unwrap(); + + let centroids_flat = kmeans.centroids.as_primitive::().clone(); + let centroids_fsl = + Arc::new(FixedSizeListArray::try_new_from_values(centroids_flat, dim as i32).unwrap()); + let mut ivf_params = + IvfBuildParams::try_with_centroids(TWO_FRAG_NUM_PARTITIONS, centroids_fsl).unwrap(); + ivf_params.max_iters = TWO_FRAG_MAX_ITERS as usize; + ivf_params.sample_rate = TWO_FRAG_SAMPLE_RATE; + ivf_params + } + + async fn build_segments_for_fragment_groups( + dataset: &mut Dataset, fragment_groups: Vec>, // each group is a set of fragment ids params: &VectorIndexParams, index_name: &str, @@ -6885,6 +7302,479 @@ mod tests { assert!(result.num_rows() > 0); } + /// A distributed (sharded, precomputed-IVF) build must support covering ("included") + /// columns: each shard writes the covered columns into its own segment's storage, the + /// segment metadata advertises them as the trailing entries of `fields`, and a covered + /// projection over the committed index skips the base-table take while returning + /// correct, row-aligned covered values. + #[tokio::test] + async fn test_distributed_vector_build_supports_covering_columns() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + let vector_field_id = dataset.schema().field("vector").unwrap().id; + let id_field_id = dataset.schema().field("id").unwrap().id; + let payload_field_id = dataset.schema().field("payload").unwrap().id; + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["id".to_string(), "payload".to_string()]); + + // One covered segment per fragment, over ALL fragments, so the committed index gives + // full coverage -- no uncovered flat delta-scan can mask a covering defect. Each shard + // reads the covering columns from its own fragment subset and writes them into its + // segment's storage. + let mut segments = Vec::new(); + for fragment in fragments.iter() { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!( + segment.covering_fields, + vec![id_field_id, payload_field_id], + "distributed covered segment must record both covered columns' field ids" + ); + // Carried columns are a trailing subset of `fields` with the keyed vector field + // first; `IndexMetadata::validate_covering_fields` rejects any other shape. + assert_eq!( + segment.fields, + vec![vector_field_id, id_field_id, payload_field_id], + "covered fields must be the trailing entries of the segment's fields" + ); + segments.push(segment); + } + + dataset + .commit_existing_index_segments("vec_idx", "vector", segments) + .await + .unwrap(); + + let query_batch = dataset + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(1), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query_batch["vector"].as_fixed_size_list().value(0); + let q = q.as_primitive::(); + + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "payload"]).unwrap(); + + // The base-table take renders as `LanceRead` in the explained plan, so that -- not + // the string "Take" -- is what its absence has to be asserted on. A flat-KNN + // fallback has no take either, so pin the ANN path first or the assertion below + // could pass without the index being read at all. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the committed index; plan was:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "covered projection ['id','payload'] should skip the base-table take; plan was:\n{plan}" + ); + + let covered = scan.try_into_batch().await.unwrap(); + let row_ids = covered + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(!row_ids.is_empty()); + + // Every covered column -- including the nullable `payload` with its nulls -- must be + // row-aligned with an independent base-table take. Full array equality checks values + // AND validity, so a dropped or shifted null slot fails here; the null count is + // asserted first so that claim cannot rest on an all-valid result set. + assert!( + covered + .column_by_name("payload") + .expect("covered 'payload' column") + .null_count() + > 0, + "the returned rows must include a null covered 'payload'" + ); + let row_id_vec: Vec = row_ids.values().to_vec(); + let base = dataset + .take_rows( + &row_id_vec, + crate::dataset::ProjectionRequest::from_columns( + ["id", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + for col in ["id", "payload"] { + assert_eq!( + covered.column_by_name(col).expect("covered column"), + base.column_by_name(col).expect("base-table column"), + "covered '{col}' must match the base-table value (incl. nulls) for each row" + ); + } + + // Take elision is worthless if the covered search returns the wrong neighbors, so also + // gate on recall against brute-force ground truth (all 4 partitions probed). + let returned: HashSet = row_id_vec.iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered distributed build recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// The cross-shard segment *merger* must also carry covering columns: merging covered + /// shards into one unified auxiliary index must retain the covered payload (not drop it at + /// the merger's rebuilt output schema), so a covered projection over the merged index still + /// skips the take and returns correct values. Parametrized over quantizer family so each + /// type's internal-name list (which the merger uses to detect the covering columns) is + /// exercised on the covered-merge path, including IVF_HNSW_PQ for the HNSW-variant lists. + // Note: IVF_RQ (RaBitQ) is intentionally excluded -- distributed RQ *merge* fails + // independently of covering (shards train different `fast_rotation_signs`, so the merger's + // structural-equality check rejects them); covered RQ works on the per-segment-commit path. + /// The HNSW arms build with `lightweight_hnsw_params()`: this test asserts that the merger + /// carries covering columns through each quantizer's internal-name list, which graph quality + /// has no bearing on, and a full-size graph would push the cases past the one-second budget. + /// `hnsw_flat` and `hnsw_sq` are not redundant with `flat`/`sq`: they are the only cases that + /// reach the `IvfHnswFlat` / `IvfHnswSq` arms of `covering_fields_from_shard_schema`. Without + /// them, mapping `IvfHnswSq` to `FLAT_INTERNAL_COLUMNS` would classify `__sq_code` as a + /// covering column -- a duplicate field in the writer schema -- with no test failing. + #[rstest] + #[case::pq("IVF_PQ")] + #[case::sq("IVF_SQ")] + #[case::flat("IVF_FLAT")] + #[case::hnsw_pq("IVF_HNSW_PQ")] + #[case::hnsw_flat("IVF_HNSW_FLAT")] + #[case::hnsw_sq("IVF_HNSW_SQ")] + #[tokio::test] + async fn test_distributed_vector_merge_supports_covering_columns(#[case] index_type: &str) { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered_merge_{index_type}", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + let id_field_id = dataset.schema().field("id").unwrap().id; + let payload_field_id = dataset.schema().field("payload").unwrap().id; + + let mut params = match index_type { + "IVF_PQ" => { + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params) + } + "IVF_SQ" => VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + SQBuildParams::default(), + ), + "IVF_FLAT" => VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + ), + "IVF_HNSW_PQ" => { + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + VectorIndexParams::with_ivf_hnsw_pq_params( + DistanceType::L2, + ivf_params, + lightweight_hnsw_params(), + pq_params, + ) + } + "IVF_HNSW_FLAT" => VectorIndexParams::ivf_hnsw( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + lightweight_hnsw_params(), + ), + "IVF_HNSW_SQ" => VectorIndexParams::with_ivf_hnsw_sq_params( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + lightweight_hnsw_params(), + SQBuildParams::default(), + ), + other => panic!("unexpected index type {other}"), + }; + params.covering_columns(vec!["id".to_string(), "payload".to_string()]); + + // All fragments, so the merged index gives full coverage (no uncovered delta scan). + let mut segments = Vec::new(); + for fragment in fragments.iter() { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + segments.push(segment); + } + + // Merge the covered shards into one unified segment, then commit it. + let merged = dataset + .merge_existing_index_segments(segments) + .await + .unwrap(); + assert_eq!( + merged.covering_fields, + vec![id_field_id, payload_field_id], + "merged covered segment must retain both covered columns' field ids" + ); + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![merged]) + .await + .unwrap(); + + let query_batch = dataset + .scan() + .project(&["vector"] as &[&str]) + .unwrap() + .limit(Some(1), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query_batch["vector"].as_fixed_size_list().value(0); + let q = q.as_primitive::(); + + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(4); + scan.with_row_id(); + scan.project(&["id", "payload"]).unwrap(); + + // As above: pin the ANN path, then assert the take (`LanceRead`) is gone. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the merged index; plan was:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "covered projection over the merged index should skip the base-table take; \ + plan was:\n{plan}" + ); + + let covered = scan.try_into_batch().await.unwrap(); + let row_ids = covered + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::(); + assert!(!row_ids.is_empty()); + + // Both covered columns (incl. the nullable `payload` with its nulls) must be + // row-aligned with an independent base-table take over the merged index. Full array + // equality checks validity too, and the null count is asserted first so that claim + // cannot rest on an all-valid result set. + assert!( + covered + .column_by_name("payload") + .expect("covered 'payload' column") + .null_count() + > 0, + "the returned rows must include a null covered 'payload'" + ); + let row_id_vec: Vec = row_ids.values().to_vec(); + let base = dataset + .take_rows( + &row_id_vec, + crate::dataset::ProjectionRequest::from_columns( + ["id", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + for col in ["id", "payload"] { + assert_eq!( + covered.column_by_name(col).expect("covered column"), + base.column_by_name(col).expect("base-table column"), + "covered '{col}' over the merged index must match the base-table value for each row" + ); + } + + // Gate on recall too, so a merge that returns the wrong neighbors is caught. + let returned: HashSet = row_id_vec.iter().copied().collect(); + let truth = ground_truth(&dataset, "vector", q, 10, DistanceType::L2).await; + let recall = truth.intersection(&returned).count() as f32 / truth.len() as f32; + assert!( + recall >= 0.5, + "covered merged {index_type} recall {recall} < 0.5 (returned {returned:?}, truth {truth:?})" + ); + } + + /// A merged covered segment must itself be usable as an input to a later merge -- + /// the hierarchical/incremental distributed workflow. The merger requires every + /// covered shard to carry its source field ids, so the merged output has to carry + /// them too; otherwise the merger rejects its own product and tells the user to + /// rebuild a shard that only the merger can produce. + #[tokio::test] + async fn test_distributed_covered_merge_output_can_be_merged_again() { + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let dataset_uri = format!("{}/distributed_covered_remerge", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let fragments = dataset.get_fragments(); + assert!( + fragments.len() >= 3, + "need >= 3 fragments so the second-stage merge has more than one input" + ); + // IVF_FLAT: a merged *PQ* segment stores transposed codes and the merger requires + // row-major shards, so PQ cannot be re-merged for reasons unrelated to covering. + // FLAT and SQ can, which is where the missing stamp actually bit. + let mut params = VectorIndexParams::with_ivf_flat_params( + DistanceType::L2, + prepare_covered_ivf(&dataset, "vector").await, + ); + params.covering_columns(vec!["payload".to_string()]); + + let mut segments = Vec::new(); + for fragment in fragments.iter() { + segments.push( + dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + + // Stage 1: merge the first two shards. + let rest = segments.split_off(2); + let merged = dataset + .merge_existing_index_segments(segments) + .await + .unwrap(); + // Stage 2: feed that merged segment back in alongside the remaining shards. This + // is the hierarchical workflow the missing stamp used to make impossible. + let mut stage_two = vec![merged]; + stage_two.extend(rest); + let remerged = dataset + .merge_existing_index_segments(stage_two) + .await + .expect("a merged covered segment must be mergeable again"); + assert_eq!( + remerged.covering_fields, + vec![dataset.schema().field("payload").unwrap().id], + "the re-merged segment must still declare the covered field" + ); + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![remerged]) + .await + .unwrap(); + } + + /// Same covering column *name and type* across shards, but two different logical + /// fields: `payload` is dropped and re-added between the two shard builds, so it + /// comes back with a fresh field id. Name+type comparison accepts this and + /// `concat_batches` would then stack shard A's values for the dropped field under + /// the re-added one -- covered queries over shard A's rows serving the old column. + /// + /// Carried columns live in `fields`, so the two shards disagree on `fields` as well as + /// on `covering_fields`, and the segment-level check rejects them before the auxiliary + /// merger runs. The merger's own `COVERING_FIELD_IDS_KEY` comparison guards the same + /// case one layer down, for drivers calling `merge_partial_vector_auxiliary_files` + /// directly; what this test pins is that the dataset API never lets the corruption + /// through. + #[tokio::test] + async fn test_distributed_vector_merge_rejects_covering_of_different_field_ids() { + use crate::dataset::NewColumnTransform; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + + let payload = Arc::new(UInt64Array::from_iter_values(0..TWO_FRAG_NUM_ROWS as u64)); + let values = generate_random_array_with_range(TWO_FRAG_NUM_ROWS * COVERED_DIM, 0.0..1.0); + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(values), COVERED_DIM as i32) + .unwrap(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::UInt64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![payload, vectors]).unwrap(); + let dataset_uri = format!("{}/distributed_covered_merge_field_id", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, vec![batch]).await; + + let fragments = dataset.get_fragments(); + assert!(fragments.len() >= 2); + + let (ivf_params, pq_params) = prepare_covered_ivf_pq(&dataset, "vector").await; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::L2, ivf_params, pq_params); + params.covering_columns(vec!["payload".to_string()]); + + // Shard A covers `payload` as it exists now. + let segment_a = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragments[0].id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + let original_field_id = dataset.schema().field("payload").unwrap().id; + + // Drop and re-add `payload`: same name, same type, new field id. The re-add is + // metadata-only (AllNulls), so no data file changes and nothing else notices. + dataset.drop_columns(&["payload"]).await.unwrap(); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(arrow_schema::Schema::new(vec![ + Field::new("payload", DataType::UInt64, true), + ]))), + None, + None, + ) + .await + .unwrap(); + let new_field_id = dataset.schema().field("payload").unwrap().id; + assert_ne!( + original_field_id, new_field_id, + "re-adding the column must produce a new field id, or this test proves nothing" + ); + + let fragments = dataset.get_fragments(); + let segment_b = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(vec![fragments[1].id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + + let err = dataset + .merge_existing_index_segments(vec![segment_a, segment_b]) + .await + .expect_err("merging shards covering different field ids must fail"); + let msg = err.to_string(); + assert!( + msg.contains("identical fields"), + "error should describe the covered-field disagreement; was: {msg}" + ); + } + #[rstest] #[case::flat("IVF_HNSW_FLAT")] #[case::pq("IVF_HNSW_PQ")]