From a17d7cea537e9e92f652ed10160109736d15a7d8 Mon Sep 17 00:00:00 2001 From: Vivek Date: Sun, 16 Aug 2026 16:11:42 -0700 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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")] From ecaa23157c7a49ce2736caf182bc14636df3c16f Mon Sep 17 00:00:00 2001 From: Vivek Date: Fri, 28 Aug 2026 11:30:38 -0700 Subject: [PATCH 5/7] feat(format): specify carried-column storage and allow carrying a keyed column #8535 declared `IndexMetadata.covering_fields` but left the physical side unspecified; this specifies it -- carried values are extra columns in `auxiliary.idx`, discovered by exclusion against the quantizer's internal columns and bound to their source fields by a new `covering_field_ids` metadata key, with no `index_version` bump. It also permits an index to carry a column it is also keyed on, the only case where an id repeats in `fields`, so readers must take the carried set from `covering_fields` rather than subtract the keyed prefix. The "Current state" note becomes a rule about verifying each segment rather than a snapshot of which writers exist, so it stays accurate as implementations land instead of needing an edit to this spec each time one does. Adds `VectorQueryProto.covering_projection` (field 15) to reserve the tag, with the one initializer the new field forces on `query_to_proto`; no writer emits carried values yet, so the implementation follows separately. --- docs/src/format/index/index.md | 57 ++++++++++++++++++--------- docs/src/format/index/vector/index.md | 24 +++++++++++ protos/ann.proto | 23 +++++++++++ rust/lance/src/io/exec/ann_proto.rs | 3 ++ 4 files changed, 88 insertions(+), 19 deletions(-) diff --git a/docs/src/format/index/index.md b/docs/src/format/index/index.md index 01970133208..6121011e750 100644 --- a/docs/src/format/index/index.md +++ b/docs/src/format/index/index.md @@ -100,10 +100,11 @@ Index segments are created and updated through a transactional process: - `name`: The index name (must match existing segments if adding to an existing index) - `fields`: The columns the index depends on: the keyed column(s) it is searched on, followed by any merely-carried columns named in `covering_fields`. `fields[0]` is always a keyed column. - - `covering_fields`: The trailing subset of `fields` whose values the index carries but is not - keyed on, letting a query that only projects those columns be answered without a fragment take. - Empty for an index that carries no extra columns. Declaring a column here does not by itself - make it servable -- see [Serving carried columns](#serving-carried-columns). + - `covering_fields`: The trailing subset of `fields` whose values the index carries, letting a + query that only projects those columns be answered without a fragment take. Usually these are + columns the index is not keyed on, but a keyed column may also be carried. Empty for an index + that carries no extra columns. Declaring a column here does not by itself make it servable -- + see [Serving carried columns](#serving-carried-columns). - `fragment_bitmap`: The set of fragment IDs covered by this segment - `index_details`: Index-specific configuration and parameters - `version`: The format version of this index type @@ -141,18 +142,36 @@ fragments that would have been covered by that segment. carries. It does not establish that the segment's storage holds their values. **The segment's storage schema is authoritative.** Before answering a query from a -carried column, an engine must confirm that column is present in the storage it opened, -and fall back to a take against the base table when it is not. A segment whose -declaration names a column its storage does not hold is a legal state, not corruption: -a maintenance operation that cannot carry the payload through a rebuild is permitted to -withdraw it and leave the declaration standing. - -!!! note "Current state" - - No index builder writes carried values yet, so today every declaration is ahead of - its storage. Engines that read `covering_fields` must therefore treat it purely as a - declaration and serve every column from the base table until they have verified the - storage themselves. This is transitional; the rule above is not. +carried column, an engine must confirm that column is present and bound to the declared +logical field in the storage it opened, and fall back to a take against the base table +when it cannot. A segment whose declaration names a column its storage does not hold is +a legal state, not corruption: a maintenance operation that cannot carry the payload +through a rebuild is permitted to withdraw it and leave the declaration standing. + +!!! note "Capability varies by segment" + + Whether a segment's storage holds a declared column depends on the index type, on the + writer that produced the segment, and on what later maintenance did to it, so one + logical index may hold values for some of its segments and not others. An engine + therefore verifies each selected segment rather than inferring capability from the + index type, the writer version, or the declaration alone, and serves from the base + table every column it cannot verify. + +!!! note "A keyed column may also be carried" + + `covering_fields` usually names columns the index is *not* keyed on, but an index is + permitted to carry a column it is also keyed on -- for instance a vector index that + keeps full-precision vectors so a refine pass can re-rank without a base-table take. + The id then appears twice in `fields`, once as `fields[0]` and again as the trailing + carried entry, and once in `covering_fields`. This is the only case in which an id + repeats in `fields`. + + A reader must therefore take the carried set from `covering_fields` directly, and + never derive it by subtracting the keyed prefix from `fields`: that set difference + silently drops a column that is both. The trailing-subset rule is stated over + `covering_fields` and is unaffected: `fields[0]` remains the column the index is + searched on, and an engine serves the repeated id from storage like any other + carried column. ## Loading an index @@ -175,9 +194,9 @@ The `IndexMetadata` message contains important information about the index segme - `fields`: the columns the index depends on: the keyed column(s) the index is searched on, followed by any columns it merely carries, as named in `covering_fields`. `fields[0]` is always a keyed column. - `covering_fields`: the trailing subset of `fields` whose values the index carries alongside its own - data but is not keyed on. Empty for an index that carries no extra columns. This declaration is - not authoritative for what the segment can serve -- see - [Serving carried columns](#serving-carried-columns). + data -- usually columns it is not keyed on, though a keyed column may also be carried. Empty for an + index that carries no extra columns. This declaration is not authoritative for what the segment can + serve -- see [Serving carried columns](#serving-carried-columns). - `fragment_bitmap`: the set of fragment IDs covered by this index segment. - `index_details`: a protobuf `Any` message that contains index-specific details, such as index type, parameters, and storage format. This allows different index types to store their own metadata. diff --git a/docs/src/format/index/vector/index.md b/docs/src/format/index/vector/index.md index 91d73a23a16..709b9a58cc7 100644 --- a/docs/src/format/index/vector/index.md +++ b/docs/src/format/index/vector/index.md @@ -161,6 +161,19 @@ the Arrow schema of the Lance file varies depending on the quantization method u !!! note All partitions are stored in the same file, and partitions must be written in order. +Every quantization format below lists only its internal columns. When a V3 IVF +writer materializes carried values, it appends one trailing column per carried +field after them, named and typed exactly as in the dataset schema. This physical +payload may be a subset of the manifest's `covering_fields` declaration (see +[Index Metadata](../index.md)). A reader returns only columns whose physical +schema and source field ids it verifies across every selected segment; all other +projected columns come from a base-table take. + +A reader discovers carried columns by exclusion, not by position: any column in the +auxiliary file's schema that is not one of the quantizer's internal columns is a +carried column. Writers append them in trailing order, but a reader must not depend +on that ordering to identify them. + ##### FLAT No quantization applied - stores original vectors in their full precision: @@ -229,6 +242,17 @@ Contains RabitQ-specific metadata in JSON format (only present for RQ quantizati This includes the rotation matrix position, number of bits, and packing information. See the RQ metadata specification in the "storage_metadata" section below. +##### "covering_field_ids" + +The *source dataset* field ids of the storage file's physical carried columns, +comma separated in physical schema order (only present when the storage carries +values). Arrow fields carry no Lance field id, so names and types alone cannot +prove which logical column a payload came from. Readers use these ids to bind +physical values to the segment's `covering_fields` declaration, and treat missing, +malformed, ambiguous, or mismatched metadata as no servable carried capability. +Distributed merges use the same identity to reject shards whose columns match by +name and type but come from different fields. + ##### "storage_metadata" Contains quantizer-specific metadata as a list of JSON strings. diff --git a/protos/ann.proto b/protos/ann.proto index f5de5e25e7b..fd91c3c6d27 100644 --- a/protos/ann.proto +++ b/protos/ann.proto @@ -23,6 +23,23 @@ enum VectorApproxMode { Accurate = 2; } +// The covering ("included") index columns a query needs materialized, by name. +// +// Exists as a message rather than a bare `repeated string` because proto3 has no +// presence tracking for repeated fields, and the empty list is a distinct, meaningful +// state here: +// +// * absent: no narrowing computed; materialize every covering column declared. +// * present and empty: materialize nothing, though the index does declare covering. +// * present and non-empty: materialize exactly these. +// +// Collapsing "present, empty" into "absent" costs no correctness (a covering column is +// semantically transparent -- its values otherwise arrive from a base-table read) but +// silently restores full materialization on every distributed plan. +message CoveringProjection { + repeated string columns = 1; +} + // Serialized vector query parameters. message VectorQueryProto { // Query vector as Arrow IPC bytes (supports Float16, Float32, Float64, UInt8, etc.) @@ -43,6 +60,12 @@ message VectorQueryProto { // Query-time approximation mode. Currently only affects RQ-quantized vector // indexes, such as IVF_RQ. Other index types ignore this setting. VectorApproxMode approx_mode = 14; + // Which covering columns the index must materialize for this query. Absent means + // "not computed" -- see CoveringProjection. Carried across the wire so a remote + // executor declares the same search output schema the planner did; without it the + // executor's node is wider than the plan it came from, and the surrounding nodes + // were built against the planner's narrower schema. + CoveringProjection covering_projection = 15; } // Serializable form of ANNIvfSubIndexExec — the IVF sub-index search node. diff --git a/rust/lance/src/io/exec/ann_proto.rs b/rust/lance/src/io/exec/ann_proto.rs index c57ad4ca1b7..7ecb14211e6 100644 --- a/rust/lance/src/io/exec/ann_proto.rs +++ b/rust/lance/src/io/exec/ann_proto.rs @@ -118,6 +118,9 @@ pub fn query_to_proto(query: &Query) -> Result { dist_q_c: Some(query.dist_q_c), query_parallelism: Some(query.query_parallelism), approx_mode: approx_mode_to_proto(query.approx_mode) as i32, + // No planner narrows the covering projection yet, so this is always absent: + // "materialize every covering column declared". See `CoveringProjection`. + covering_projection: None, }) } From 3b1b34ce30604bf290c7d9fbbc1efcb1323c6386 Mon Sep 17 00:00:00 2001 From: Vivek Date: Wed, 26 Aug 2026 10:35:54 -0700 Subject: [PATCH 6/7] perf(index)!: read only the covering values a query actually needs A covered index materialised every column it carries on every query and loaded those columns with every partition it probed, caching them beside the quantization codes -- so a query touching none of them still paid for all of them twice, and a wide covering column evicted the codes it shared a cache entry with. This narrows both reads: the scanner declares which covering columns a query actually reads, and partition loads now fetch only the storage's own columns, with covering values read by position for the rows that survive scoring and not cached. BREAKING CHANGE: `lance_index::vector::Query` gains a `covering_projection` field. Construct it with `..Default::default()`, or pass `None` for the previous behaviour. BREAKING CHANGE: `IvfPq::load_partition_storage` gains a required `columns: PartitionColumns` parameter, between `partition_id` and `io_stats`. Partition loads now read only the storage's own internal columns by default, so a caller must say which set it wants: `PartitionColumns::Internal` reproduces the new default and `PartitionColumns::All` reproduces the previous behaviour of loading the covering columns too. --- docs/src/guide/performance.md | 14 + java/lance-jni/src/utils.rs | 10 + python/src/dataset.rs | 4 + rust/lance-index/src/vector.rs | 57 + rust/lance-index/src/vector/storage.rs | 464 ++- .../src/transaction/index_maintenance.rs | 62 + .../vector/hnsw/mem_wal_recall_hnsw.rs | 1 + rust/lance/src/dataset/scanner.rs | 708 ++++- rust/lance/src/dataset/tests/dataset_index.rs | 912 ++++++ .../src/dataset/tests/dataset_scanner.rs | 1 + rust/lance/src/index.rs | 1 + rust/lance/src/index/append.rs | 51 +- rust/lance/src/index/covering.rs | 283 ++ rust/lance/src/index/vector.rs | 197 +- rust/lance/src/index/vector/builder.rs | 72 +- rust/lance/src/index/vector/fixture_test.rs | 1 + rust/lance/src/index/vector/hamming.rs | 5 +- rust/lance/src/index/vector/ivf.rs | 7 +- rust/lance/src/index/vector/ivf/v2.rs | 2566 +++++++++++++++-- rust/lance/src/index/vector/pq.rs | 1 + rust/lance/src/index/vector/utils.rs | 37 +- rust/lance/src/io/exec/ann_proto.rs | 217 +- rust/lance/src/io/exec/knn.rs | 410 ++- 23 files changed, 5691 insertions(+), 390 deletions(-) create mode 100644 rust/lance/src/index/covering.rs diff --git a/docs/src/guide/performance.md b/docs/src/guide/performance.md index 8d7f0c5401f..e749f68a1d0 100644 --- a/docs/src/guide/performance.md +++ b/docs/src/guide/performance.md @@ -529,3 +529,17 @@ Set `LANCE_DISABLE_AMX=1` to take the AMX paths out of service without rebuildin A/B measurement, or to get the previous behaviour back. Because it also moves partition assignment back to the approximate path, an index built with it set is not equivalent to one built without it; compare recall, not just build time. + +#### Covering Columns + +A vector index can carry the values of extra columns beside its vectors, so a query whose +projection they satisfy is answered from the index without a take against the base table. +That trade only pays off where the search settles into a single global top-k heap, because +the covering read is then bounded by the query's survivors. Every HNSW index, and any query +with `query_parallelism` above one, emits results per partition instead, so covering reads +scale with the number of partitions probed rather than with `k`, and are issued serially. + +On those shapes covering is *slower* than the base-table take it exists to avoid — roughly +2.8x a plain index's latency warm at nprobe 32, and 1.18x cold. Prefer IVF_PQ or IVF_FLAT +when declaring covering columns; index creation logs a warning when covering is combined +with HNSW. diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 09bc09207da..29feff8f419 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -261,6 +261,11 @@ pub fn build_compaction_options( } // Convert from Java Optional to Rust Option +// +// This builds a `Query` directly rather than through `Scanner::nearest`, and is used +// only by `JniTestHelper.parseQuery` to check the Java-side field marshalling. Java's +// real search path is `blocking_scanner.rs`, which configures a `Scanner`, so it picks +// up every plan-derived setting -- including the covering projection below. pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> { let query = env.get_optional(&query_obj, |env, java_obj| { let column = env.get_string_from_method(&java_obj, "getColumn")?; @@ -305,6 +310,11 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> dist_q_c: 0.0, query_parallelism, approx_mode, + // Not a user-settable search parameter: the covering projection is derived + // per plan from what the scan reads, so it stays `None` (materialize + // whatever the index declares) at the binding boundary. Java's real search + // path resolves it in `Scanner`; see the note on this function. + covering_projection: None, }) })?; diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 547ba91d16b..ffa00a78a87 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -5674,6 +5674,10 @@ impl PySearchFilter { query_parallelism, dist_q_c: 0.0, approx_mode, + // Not a user-settable search parameter: the covering projection is derived + // from what the plan reads, so it stays `None` (materialize whatever the + // index declares) at the binding boundary. + covering_projection: None, }; Ok(Self { diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index a41973f57c8..fd0f5639f6e 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -166,6 +166,52 @@ pub struct Query { /// This currently only affects RQ-quantized vector indexes, such as IVF_RQ. /// Other index types ignore this setting. pub approx_mode: ApproxMode, + + /// The covering ("included") columns this query needs the index to materialize, + /// by name. + /// + /// Unlike the other fields here this is not a user-facing search parameter. It is + /// derived per plan from what the query actually reads, so that an index covering a + /// wide payload does not pay to materialize that payload on searches which never + /// look at it. + /// + /// # The three states are distinct, and must stay distinct + /// + /// * `None` — no query-level narrowing or physical-capability resolution was computed. + /// A raw index search may consider every physical covering column, but an execution + /// planner must not treat this state as proof that declared values are servable. It + /// resolves storage first and converts the query to an explicit `Some` projection. + /// * `Some(&[])` — the index *does* declare covering columns, but this query needs + /// **none** of them. Consumers must do no covering work at all: not "project zero + /// columns out of a batch that was loaded anyway", but skip the covering read + /// entirely. + /// * `Some(cols)` — this query needs exactly `cols`, and the execution planner has + /// verified that every selected segment can physically serve them. + /// + /// `Some(&[])` is the state this field exists for, and the one that silently + /// degrades if it is folded into `None`: a covering column is semantically + /// transparent, so conflating the two restores full materialization while every + /// result stays byte-identical (the values simply arrive from the base table + /// instead). No result-correctness test can catch that; only a cost or plan-shape + /// assertion can. Treat `Option::unwrap_or_default()` on this field, or any + /// `is_empty()` test that does not first distinguish `None`, as a bug. + /// + /// # Physical capability is a ceiling + /// + /// A declared column left out of `cols` is simply fetched from the base table — + /// correct, just slower. A column listed here that any selected segment cannot emit + /// leaves the covered projection short a column. Planners may therefore remove an + /// unproven field, but must never add one based on the manifest declaration alone. + /// + /// # Notes + /// + /// Names rather than field ids, because consumers match against a storage batch whose + /// columns carry names. The planner separately verifies source field ids, Arrow shape, + /// and compatible physical order before producing this name projection. + /// + /// `Arc` rather than `Vec` because [`Query`] is cloned once per probed partition on + /// the search hot path. + pub covering_projection: Option>, } impl From for DistanceType { @@ -236,6 +282,17 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { /// Get the total number of partitions in the index. fn total_partitions(&self) -> usize; + /// Covering columns physically present and safe to serve from this index, paired with + /// their source dataset field ids and returned in storage order. + /// + /// [`IndexMetadata::covering_fields`](lance_table::format::IndexMetadata::covering_fields) + /// is only the logical declaration. Query planners intersect that declaration with this + /// physical capability for every segment before omitting a base-table read. Formats that + /// do not expose a verifiable covering payload use the default empty capability. + fn physical_covering_fields(&self) -> Result> { + Ok(Vec::new()) + } + /// Search a single partition for nearest neighbors. /// /// This method should return the same results as [`VectorIndex::search`] method except diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 493e28d200d..e57da92df56 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -5,14 +5,14 @@ use crate::vector::quantizer::QuantizerStorage; use arrow::compute::concat_batches; -use arrow_array::{ArrayRef, RecordBatch}; -use arrow_schema::SchemaRef; +use arrow_array::{ArrayRef, RecordBatch, UInt32Array}; +use arrow_schema::{Field, Schema, SchemaRef}; use futures::prelude::stream::TryStreamExt; use lance_arrow::RecordBatchExt; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, ROW_ID, Result}; use lance_encoding::decoder::FilterExpression; -use lance_file::reader::FileReader; +use lance_file::reader::{FileReader, ReaderProjection}; use lance_io::ReadBatchParams; use lance_io::scheduler::IoStats; use lance_linalg::distance::DistanceType; @@ -20,7 +20,7 @@ use prost::Message; use std::{ any::Any, borrow::Cow, - collections::BinaryHeap, + collections::{BinaryHeap, HashSet}, mem::size_of, ops::{Deref, DerefMut}, sync::Arc, @@ -156,7 +156,7 @@ 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. +/// covering ("included") columns, comma separated in physical schema 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 @@ -405,6 +405,73 @@ pub(crate) fn covering_field_indices_excluding( .collect() } +/// Resolve the covering columns a physical storage schema can safely serve. +/// +/// Arrow fields do not carry Lance field ids, so the physical columns are usable only +/// when the storage also stamps one source id for each column. Invalid or incomplete +/// capability metadata deliberately resolves to no columns: the manifest declaration is +/// still valid, but readers must fetch those values from the base table instead. +fn physical_covering_fields_from_schema(schema: &Schema, internal: &[&str]) -> Vec<(i32, Field)> { + let covering_indices = covering_field_indices_excluding(schema, internal); + if covering_indices.is_empty() { + return Vec::new(); + } + + let Some(encoded_ids) = schema.metadata().get(COVERING_FIELD_IDS_KEY) else { + // Every return of `Vec::new()` below withdraws the whole covering payload, and the + // query then reads those columns from the base table with correct results and no + // error -- so without a line here the operator cannot tell a covered index that + // stopped eliding takes from one that never did. + log::debug!( + "Index storage carries {} covering column(s) but no `{}` stamp, so none can be \ + served; their values will come from a base-table read.", + covering_indices.len(), + COVERING_FIELD_IDS_KEY + ); + return Vec::new(); + }; + let Ok(source_ids) = encoded_ids + .split(',') + .map(str::trim) + .map(str::parse::) + .collect::, _>>() + else { + return Vec::new(); + }; + if source_ids.len() != covering_indices.len() + || source_ids.iter().collect::>().len() != source_ids.len() + { + log::debug!( + "Index storage's `{}` stamp ({:?}) does not pair one distinct source id with each \ + of its {} covering column(s), so none can be served; their values will come from \ + a base-table read.", + COVERING_FIELD_IDS_KEY, + source_ids, + covering_indices.len() + ); + return Vec::new(); + } + + let covering_names = covering_indices + .iter() + .map(|index| schema.field(*index).name()) + .collect::>(); + if covering_names.len() != covering_indices.len() { + log::debug!( + "Index storage has duplicate covering column names, so a value cannot be bound to \ + one source id; none can be served and their values will come from a base-table \ + read." + ); + return Vec::new(); + } + + source_ids + .into_iter() + .zip(covering_indices) + .map(|(source_id, index)| (source_id, schema.field(index).clone())) + .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` @@ -637,6 +704,61 @@ impl StorageBuilder { } } +/// How [`IvfQuantizationStorage::take_covering`] reads a partition's covering values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CoveringRead { + /// [`ReadBatchParams::Indices`] -- read only the survivors' rows. + Scattered, + /// [`ReadBatchParams::Range`] -- read the partition's whole row range. + Sequential, +} + +/// Survivors at or above this percentage of a partition's rows make the scattered read +/// degenerate: `Indices` over most of a range costs more than the sequential read of that +/// range it replaced, because the scattered form gives up contiguity while still touching +/// nearly every page. +/// +/// # Derivation +/// +/// Measured, not guessed. At production partition sizes (~8,192 rows) and `k <= 100` the +/// survivors of one probe are **<= 1.2%** of a partition; on the 2,048-row fixture +/// partitions with `k = 10` it is 0.49%. Ten percent therefore sits roughly an order of +/// magnitude above the worst observed common case, so the fallback costs nothing in +/// practice while still bounding the degenerate case at "no worse than the read this +/// replaced". The design's §1 warns explicitly against assuming random access always +/// wins; this is that guard, deliberately conservative in the direction that cannot +/// regress. +pub const COVERING_SCATTERED_READ_MAX_PERCENT: usize = 10; + +/// Which read shape to use for `survivors` covering rows out of `partition_rows`. +/// +/// A zero-row partition has nothing to read either way; it is reported as sequential so +/// callers have a single empty-range path. +pub fn covering_read_for(survivors: usize, partition_rows: usize) -> CoveringRead { + if partition_rows == 0 + || survivors.saturating_mul(100) + >= partition_rows.saturating_mul(COVERING_SCATTERED_READ_MAX_PERCENT) + { + CoveringRead::Sequential + } else { + CoveringRead::Scattered + } +} + +/// Which columns a partition load reads out of the storage file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PartitionColumns { + /// Only the storage's own columns ([`VectorStore::INTERNAL_COLUMNS`]): the row id, + /// the quantization codes and any per-row factor columns. Covering columns are left + /// unread, which is what makes the loaded partition independent of the query that + /// loaded it -- and therefore sound to share under a partition-id-only cache key. + Internal, + /// Every column the file carries, covering columns included. Required by the paths + /// that rewrite a partition into a new index file: a covering column that was never + /// read cannot be written back. + All, +} + /// Loader to load partitioned PQ storage from disk. #[derive(Debug)] pub struct IvfQuantizationStorage { @@ -654,6 +776,10 @@ pub struct IvfQuantizationStorage { /// decoded), but the legacy `codebook_tensor` path does decode, so this is not free /// for every index. covering_schema: std::sync::OnceLock>, + /// Lazily-computed internal-column projection (see [`Self::internal_projection`]). + /// Fixed for the storage's lifetime and consulted on every partition load, so it is + /// resolved once rather than per probe. + internal_projection: std::sync::OnceLock, } impl DeepSizeOf for IvfQuantizationStorage { @@ -726,6 +852,7 @@ impl IvfQuantizationStorage { ivf, frag_reuse_index, covering_schema: std::sync::OnceLock::new(), + internal_projection: std::sync::OnceLock::new(), }) } @@ -758,6 +885,7 @@ impl IvfQuantizationStorage { ivf, frag_reuse_index, covering_schema: std::sync::OnceLock::new(), + internal_projection: std::sync::OnceLock::new(), } } @@ -819,13 +947,55 @@ impl IvfQuantizationStorage { Ok(self.covering_schema.get_or_init(|| computed).clone()) } + /// Physical covering columns this storage can safely serve, paired with their source + /// dataset field ids and returned in storage order. + /// + /// A missing or malformed source-id stamp makes the physical covering payload + /// unavailable rather than making the index unreadable. Query planning will then use + /// the ordinary base-table take for the declared columns. + pub fn physical_covering_fields(&self) -> Vec<(i32, Field)> { + let schema = Schema::from(self.reader.schema().as_ref()); + physical_covering_fields_from_schema(&schema, Q::Storage::INTERNAL_COLUMNS) + } + /// Get the number of partitions in the storage. pub fn num_partitions(&self) -> usize { self.ivf.num_partitions() } - /// Load a partition's quantization storage, optionally measuring the exact - /// I/O it performs into `io_stats`. + /// The read projection selecting only this storage's own columns. + /// + /// [`VectorStore::INTERNAL_COLUMNS`] is the union over every build a storage type + /// has ever produced -- the legacy `__ivf_part_id`, RaBitQ's factor columns that + /// exist only above a `num_bits` threshold -- so it is intersected with this file's + /// schema first. Projecting the const directly would reject any file that legitimately + /// omits one of those columns, which is most of them. + fn internal_projection(&self) -> Result { + if let Some(cached) = self.internal_projection.get() { + return Ok(cached.clone()); + } + let schema = self.reader.schema(); + let names: Vec<&str> = Q::Storage::INTERNAL_COLUMNS + .iter() + .copied() + .filter(|name| schema.field(name).is_some()) + .collect(); + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + schema, + &names, + )?; + Ok(self.internal_projection.get_or_init(|| projection).clone()) + } + + /// Load a partition's quantization storage, reading the columns named by + /// `columns` and optionally measuring the exact I/O it performs into `io_stats`. + /// + /// [`PartitionColumns::Internal`] leaves the covering columns on disk. Search + /// wants that: covering is needed for at most `k` survivors while the codes are + /// scanned for every row on every probe, and an entry that holds no covering is + /// the same entry for every query -- which is what makes caching it under a + /// partition id alone correct rather than accidentally correct. /// /// When `io_stats` is `Some`, the partition is read through a reader whose /// scheduler also records into the sink (a cheap clone that shares all @@ -834,29 +1004,49 @@ impl IvfQuantizationStorage { pub async fn load_partition( &self, part_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result { + let projection = match columns { + PartitionColumns::Internal => Some(self.internal_projection()?), + PartitionColumns::All => None, + }; + // The batch schema must describe exactly what was read: `concat_batches` + // indexes each batch by this schema's field positions without comparing the + // two, so the full file schema over a projected read would index past the + // last column. + let schema: SchemaRef = Arc::new(match &projection { + Some(projection) => projection.schema.as_ref().into(), + None => self.reader.schema().as_ref().into(), + }); let range = self.ivf.row_range(part_id); let batch = if range.is_empty() { - let schema = self.reader.schema(); - let arrow_schema = arrow_schema::Schema::from(schema.as_ref()); - RecordBatch::new_empty(Arc::new(arrow_schema)) + RecordBatch::new_empty(schema) } else { let reader = match &io_stats { Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())), None => Cow::Borrowed(&self.reader), }; - let batches = reader - .read_stream( - ReadBatchParams::Range(range), - u32::MAX, - 1, - FilterExpression::no_filter(), - ) - .await? - .try_collect::>() - .await?; - let schema = Arc::new(self.reader.schema().as_ref().into()); + let params = ReadBatchParams::Range(range); + let stream = match projection { + Some(projection) => { + reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + } + None => { + reader + .read_stream(params, u32::MAX, 1, FilterExpression::no_filter()) + .await? + } + }; + let batches = stream.try_collect::>().await?; concat_batches(&schema, batches.iter())? }; Q::Storage::try_from_batch_with_remapper( @@ -866,12 +1056,242 @@ impl IvfQuantizationStorage { self.frag_reuse_index.clone(), ) } + + /// This storage's covering ("included") column names in storage order, narrowed to + /// `wanted`. + /// + /// `wanted` is [`Query::covering_projection`], whose three states are all meaningful: + /// `None` keeps every physical column, `Some(&[])` keeps none, and `Some(cols)` keeps + /// the intersection. Execution planners pass an explicit, storage-verified `Some`; + /// `None` is only the raw-index behavior before capability resolution. The result is + /// empty for an ordinary index and for a query that needs no covering column -- in + /// both cases the caller skips the covering read entirely. + /// + /// Order follows storage, never `wanted`: that is the order the gathered batch comes + /// back in. The planner verifies that this physical order is compatible with its + /// declaration-ordered output schema before enabling covering. + /// + /// [`Query::covering_projection`]: crate::vector::Query::covering_projection + pub fn covering_columns(&self, wanted: Option<&[String]>) -> Result> { + let Some(schema) = self.covering_schema()? else { + return Ok(Vec::new()); + }; + Ok(schema + .fields() + .iter() + .map(|field| field.name()) + .filter(|name| name.as_str() != ROW_ID) + .filter(|name| wanted.is_none_or(|wanted| wanted.contains(name))) + .cloned() + .collect()) + } + + /// The read projection [`Self::take_covering`] uses: `[_rowid, ]`. + fn covering_projection(&self, columns: &[String]) -> Result { + let mut names: Vec<&str> = Vec::with_capacity(columns.len() + 1); + names.push(ROW_ID); + names.extend(columns.iter().map(String::as_str)); + lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + &names, + ) + } + + /// The `[_rowid, ]` schema [`Self::take_covering`] returns. + /// + /// Callers declare their covered output schema from this rather than from + /// [`Self::covering_schema`] so that the schema they promise and the batches they + /// emit come from a single derivation. The two agree today, but they are computed + /// from different inputs (a projected file schema against an empty `Q::Storage`) and + /// a divergence would surface as a stream that mixes batch schemas. + pub fn covering_read_schema(&self, columns: &[String]) -> Result { + Ok(Arc::new( + self.covering_projection(columns)?.schema.as_ref().into(), + )) + } + + /// Read `[_rowid, ]` for one partition's search survivors. + /// + /// This is the covering read the search path pays after scoring has settled the heap, + /// so it is proportional to `k` rather than to the partition size -- which is why the + /// partition entry the search caches carries no covering at all (see + /// [`PartitionColumns::Internal`]) and why this result is deliberately **not** cached: + /// it is a different set of rows for every query. + /// + /// `positions` are offsets into the partition's row range, strictly ascending and + /// deduplicated (the reader's take path requires that). `None` means the caller could + /// not derive them -- a deferred fragment-reuse remap drops rows from the loaded + /// partition, so its row order no longer matches the file's -- and the whole range is + /// read instead, with the caller matching by row id. + /// + /// Row ids come back with the values, so the caller can align covering values to its + /// survivors **by row id** rather than by position, and a survivor whose row the + /// gather did not return is an error rather than a silent null. + pub async fn take_covering( + &self, + part_id: usize, + positions: Option<&[u32]>, + columns: &[String], + io_stats: Option, + ) -> Result { + let projection = self.covering_projection(columns)?; + // As in `load_partition`: `concat_batches` indexes each batch by this schema's + // field positions without comparing the two, so it must describe the projected + // read rather than the whole file. + let schema: SchemaRef = Arc::new(projection.schema.as_ref().into()); + let range = self.ivf.row_range(part_id); + if range.is_empty() { + return Ok(RecordBatch::new_empty(schema)); + } + // `ReadBatchParams::Indices` addresses rows with a `u32`, so a partition ending + // beyond that cannot be expressed as a scattered read at all. Falling back keeps + // such a file readable (just slower) instead of failing the query. + let addressable = range.end <= u32::MAX as usize; + let read = match positions { + Some(positions) if addressable => covering_read_for(positions.len(), range.len()), + _ => CoveringRead::Sequential, + }; + let params = match read { + CoveringRead::Sequential => ReadBatchParams::Range(range.clone()), + CoveringRead::Scattered => ReadBatchParams::Indices(UInt32Array::from_iter_values( + // `addressable` bounds the sum, but saturate anyway: a caller passing a + // position past the partition must reach the reader's own bounds check as + // an error, not overflow here. + positions + .unwrap_or_default() + .iter() + .map(|position| (range.start as u32).saturating_add(*position)), + )), + }; + let reader = match &io_stats { + Some(io_stats) => Cow::Owned(self.reader.with_io_stats(io_stats.recorder())), + None => Cow::Borrowed(&self.reader), + }; + let batches = reader + .read_stream_projected( + params, + u32::MAX, + 1, + projection, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + let batch = concat_batches(&schema, batches.iter())?; + match &self.frag_reuse_index { + // The loaded partition's row ids were remapped through the same index when it + // was built from this file, so the gathered ids must be too or nothing matches. + Some(frag_reuse_index) => remap_row_ids_by_name(batch, frag_reuse_index.as_ref()), + None => Ok(batch), + } + } } #[cfg(test)] mod tests { - use super::{QueryScratchCapacity, QueryScratchPool}; + use super::{ + COVERING_FIELD_IDS_KEY, COVERING_SCATTERED_READ_MAX_PERCENT, CoveringRead, + QueryScratchCapacity, QueryScratchPool, covering_read_for, + physical_covering_fields_from_schema, + }; + use arrow_schema::{DataType, Field, Schema}; use lance_core::deepsize::DeepSizeOf; + use std::collections::HashMap; + + const TEST_INTERNAL_COLUMNS: &[&str] = &["_rowid", "__pq_code"]; + + fn storage_schema(covering_fields: Vec, source_ids: Option<&str>) -> Schema { + let mut fields = vec![ + Field::new("_rowid", DataType::UInt64, false), + Field::new("__pq_code", DataType::Binary, false), + ]; + fields.extend(covering_fields); + let metadata = source_ids + .map(|ids| HashMap::from([(COVERING_FIELD_IDS_KEY.to_string(), ids.to_string())])) + .unwrap_or_default(); + Schema::new_with_metadata(fields, metadata) + } + + #[test] + fn test_physical_covering_fields_require_valid_source_ids() { + let covering = vec![ + Field::new("price", DataType::Int64, true), + Field::new("payload", DataType::Utf8, false), + ]; + let schema = storage_schema(covering.clone(), Some("17, 23")); + + let resolved = physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS); + + assert_eq!( + resolved, + vec![(17, covering[0].clone()), (23, covering[1].clone())] + ); + } + + #[test] + fn test_physical_covering_fields_treat_invalid_capability_as_absent() { + let covering = vec![ + Field::new("price", DataType::Int64, true), + Field::new("payload", DataType::Utf8, false), + ]; + for source_ids in [None, Some(""), Some("17"), Some("17,nope"), Some("17,17")] { + let schema = storage_schema(covering.clone(), source_ids); + assert!( + physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS).is_empty(), + "source ids {source_ids:?} must not prove a physical covering capability" + ); + } + + let duplicate_names = storage_schema( + vec![ + Field::new("price", DataType::Int64, true), + Field::new("price", DataType::Int64, true), + ], + Some("17,23"), + ); + assert!( + physical_covering_fields_from_schema(&duplicate_names, TEST_INTERNAL_COLUMNS) + .is_empty() + ); + } + + #[test] + fn test_physical_covering_fields_ignore_orphaned_source_id_metadata() { + let schema = storage_schema(Vec::new(), Some("17")); + assert!(physical_covering_fields_from_schema(&schema, TEST_INTERNAL_COLUMNS).is_empty()); + } + + /// The scattered read is the point of the survivor gather, so it must be chosen across + /// the whole range of survivor fractions a real query produces -- and abandoned once + /// `Indices` would scatter over most of the partition anyway, where the sequential read + /// it replaced is cheaper. + /// + /// The inputs are the measured ones: at production partition sizes (~8,192 rows) the + /// survivors of one probe are <= 1.2% at `k <= 100`, and 0.49% on the 2,048-row fixture + /// partitions at `k = 10`. Both must land on `Scattered` with room to spare, or the + /// threshold is not conservative -- it is in the way. + /// + /// The two rows either side of 10% are what make this a threshold test rather than a + /// "scattered usually wins" test: moving `COVERING_SCATTERED_READ_MAX_PERCENT` in + /// either direction fails one of them. + #[test] + fn test_covering_read_falls_back_to_a_sequential_read_near_the_threshold() { + assert_eq!(COVERING_SCATTERED_READ_MAX_PERCENT, 10); + + // k = 10 on a 2,048-row fixture partition: 0.49%. + assert_eq!(covering_read_for(10, 2048), CoveringRead::Scattered); + // k = 100 on an 8,192-row production partition: 1.2%, the measured worst case. + assert_eq!(covering_read_for(100, 8192), CoveringRead::Scattered); + // Either side of the threshold: 819/8192 = 9.998%, 820/8192 = 10.01%. + assert_eq!(covering_read_for(819, 8192), CoveringRead::Scattered); + assert_eq!(covering_read_for(820, 8192), CoveringRead::Sequential); + // The degenerate case the threshold exists for. + assert_eq!(covering_read_for(8192, 8192), CoveringRead::Sequential); + // An empty partition has nothing to read either way; one path, not two. + assert_eq!(covering_read_for(0, 0), CoveringRead::Sequential); + } #[test] fn test_query_scratch_pool_reuses_buffers() { diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index f96833fe8a3..d781ecb85eb 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -404,9 +404,17 @@ impl Transaction { match new_schema.field_by_id(covered_id) { None => true, Some(new) => { + // Metadata is compared for the same reason name, type and nullability are: + // this guard exists to turn a covering desync into an explicit error, and + // the read path's acceptance test (`covering_fields_match` in + // `io::exec::knn`) compares metadata too. A difference this guard ignores + // but that one does not is not a caught error -- it is a covered index + // that silently stops serving its columns and falls back to a base-table + // read, with no way for the user to learn why. old.name != new.name || old.nullable != new.nullable || old.data_type() != new.data_type() + || old.metadata != new.metadata } } } @@ -1124,6 +1132,60 @@ mod tests { assert!(indices.is_empty()); } + /// A covered column whose *only* change is field metadata must be rejected here. + /// + /// The read path intersects a declaration against each segment's physical columns with + /// `covering_fields_match`, which compares name, type, nullability AND metadata. If + /// this guard compared any less, a metadata-only edit would commit cleanly and then + /// silently withdraw covering at query time -- exactly the silent degradation the + /// guard is here to convert into an error. + #[test] + fn test_covered_field_subtree_change_notices_metadata_only_edit() { + let with_metadata = |value: Option<&str>| { + let mut field = ArrowField::new("carried", DataType::Int32, false); + if let Some(value) = value { + field = field.with_metadata( + [("unit".to_string(), value.to_string())] + .into_iter() + .collect(), + ); + } + LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("key", DataType::Int32, false), + field, + ])) + .unwrap() + }; + + let old_schema = with_metadata(Some("metres")); + let carried_id = old_schema.field("carried").unwrap().id; + + assert!( + Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(Some("feet")), + carried_id + ), + "a changed metadata value must count as a covering change" + ); + assert!( + Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(None), + carried_id + ), + "dropping metadata entirely must count as a covering change" + ); + assert!( + !Transaction::covered_field_subtree_changed( + &old_schema, + &with_metadata(Some("metres")), + carried_id + ), + "an unchanged field must not be reported as changed" + ); + } + #[test] fn test_prune_overlay_stale_fields_from_indices() { // Fragment 0 carried an overlay on field 1 committed at v5, and was diff --git a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs index 625900cf834..add484e8d07 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs +++ b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs @@ -596,6 +596,7 @@ async fn run_checkpoint( query_parallelism: 1, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; // IVFIndex::search is intentionally unimplemented (top-level does // partition-aware search); replicate the ANN exec node: pick the diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 21c93c051cb..02eac7bbb0e 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -100,6 +100,7 @@ use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; +use crate::index::covering::effective_covering; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, @@ -645,6 +646,33 @@ pub struct FilterPlan { expr_filter_plan: ExprFilterPlan, } +/// Whether an FTS query scores from the columns of the plan it is handed, rather than +/// building its own FTS plan and joining the result on `_rowid`. +/// +/// Only a `Match` does: it is the shape `FlatMatchQueryExec` / `FlatMatchFilterExec` +/// support, and those read the document column straight out of their input. +/// +/// This is the single authority for that rule. Two very different callers depend on it and +/// must not drift apart: +/// +/// * [`FilterPlan::refine_columns`], which decides what the refine pass has to *take* from +/// the base table, and +/// * [`Scanner::fts_scored_columns`], which decides what a covered vector index has to +/// *materialize* so that take is not needed. +/// +/// They are two halves of one decision. If a new FTS shape learns to score from its input +/// and only the first is updated, the take appears and the covered index stops declaring +/// the column it serves -- silently, because the results are identical either way. +fn fts_scores_from_input(query: &FtsQuery) -> bool { + match query { + FtsQuery::Match(_) => true, + FtsQuery::Phrase(_) + | FtsQuery::Boost(_) + | FtsQuery::MultiMatch(_) + | FtsQuery::Boolean(_) => false, + } +} + impl FilterPlan { pub fn new(query_filter: Option, expr_filter_plan: ExprFilterPlan) -> Self { Self { @@ -700,9 +728,11 @@ impl FilterPlan { fts_query.columns() }; - // Add refine column for match query since it supports `FlatMatchQueryExec`. - // Other fts query use join so we don't need to add refine column. - if let FtsQuery::Match(_) = &fts_query.query { + // A shape that scores from its input needs that column taken for it; + // the others join on `_rowid` and need nothing. Same rule, same + // authority, as the covering declaration -- see + // `fts_scores_from_input`. + if fts_scores_from_input(&fts_query.query) { columns.extend(cols.iter().cloned().collect::>()); } } @@ -2021,6 +2051,7 @@ impl Scanner { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }); self.nearest_query_count = query_count; self.is_batch_nearest = is_batch_nearest; @@ -3779,9 +3810,17 @@ impl Scanner { if self.prefilter { let source: Arc = match &filter_plan.vector_filter() { Some(vector_query) => { - // Perform vector search first then rerank according to BM25 scores + // Perform vector search first then rerank according to BM25 scores. + // The rerank reads the FTS document column from this plan, so the + // vector search must be told to materialize it when the index covers + // it -- the scan's own projection need not mention it at all. + let vector_query = self.with_covering_projection( + vector_query, + filter_plan, + &Self::fts_scored_columns(query), + ); let vector_plan = self - .vector_search(&filter_plan.expr_filter_plan, vector_query) + .vector_search(&filter_plan.expr_filter_plan, &vector_query) .await?; self.fts_rerank(vector_plan, query).await? } @@ -3810,6 +3849,15 @@ impl Scanner { let Some(query) = self.nearest.as_ref() else { return Err(Error::invalid_input("No nearest query".to_string())); }; + // A postfiltered FTS *query filter* is scored by the refine pass on top of this + // search (`FilterPlan::refine_columns` takes its document column), so name it here + // too. Free in the prefilter branch below: that branch plans the FTS index as the + // source and scores with a flat KNN, so it builds no ANN node to declare anything. + let fts_scored = filter_plan + .fts_filter() + .map(|fts| Self::fts_scored_columns(&fts)) + .unwrap_or_default(); + let query = self.with_covering_projection(query, filter_plan, &fts_scored); if self.prefilter { log::trace!("source is a vector search (prefilter)"); @@ -3826,10 +3874,10 @@ impl Scanner { .union_column(&query.column, OnMissing::Error)?; let plan = self.take(fts_plan, projection)?; - self.flat_knn(plan, query)? + self.flat_knn(plan, &query)? } None => { - self.vector_search(&filter_plan.expr_filter_plan, query) + self.vector_search(&filter_plan.expr_filter_plan, &query) .await? } }; @@ -3841,7 +3889,7 @@ impl Scanner { // If we are postfiltering then we can't use scalar indices for the filter // and will need to run the postfilter in memory filter_plan.make_refine_only(); - self.vector_search(&ExprFilterPlan::default(), query).await + self.vector_search(&ExprFilterPlan::default(), &query).await } } @@ -5258,7 +5306,187 @@ impl Scanner { Ok(flat_match_plan) } + /// The dataset columns this scan still has to read once the vector search has + /// produced its `[_distance, _rowid]` rows. + /// + /// This is the input to [`Query::covering_projection`]. A covering column only pays + /// for itself if it saves a base-table read, so the set an index should materialize + /// is this set intersected with what the index declares. The intersection (and the + /// emission order, which is the index's, not this one's) is done by + /// [`effective_covering`]; this method only answers "what does the rest of the plan + /// read". + /// + /// Every column named here is one a later [`Scanner::take`] would otherwise fetch: + /// the output projection (`create_plan`'s final take), the sort keys and the + /// aggregate inputs, plus whatever the caller passes in `also_read`. Returning a + /// *superset* is harmless -- a name the index does not declare is dropped by the + /// intersection -- so this errs wide on purpose. + /// + /// Names are top-level only, matching the covering declaration: `covering_columns` + /// rejects dotted paths, so a projection of `nested.x` contributes `nested` and a + /// filter on `nested.x` contributes nothing. Both merely under-narrow. + /// + /// `also_read` carries the columns a node *between* the search and the final + /// projection consumes, which this method cannot see from the scan options alone: + /// + /// * on a postfiltered scan, every column the filter mentions -- not just the refine + /// ones, because an index-satisfied conjunct can still be rechecked, and the recheck + /// reads from whatever the search node emitted. A *prefiltered* scan contributes + /// none of them; see [`Self::with_covering_projection`] for why; + /// * the document column an FTS scorer stacked on the vector search reads from it + /// (see [`Self::fts_scored_columns`]); it takes the column from the base table when + /// the search did not emit it. + /// + /// # An empty result is meaningful + /// + /// A query selecting only `_distance` / `_rowid` reads no dataset column at all, and + /// must yield an empty vec, which becomes `Some(&[])` -- "this query needs no + /// covering column". That is the state the whole projection pushdown exists for; + /// see [`Query::covering_projection`] for why it must not collapse into `None`. + /// + /// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection + /// [`effective_covering`]: crate::index::covering::effective_covering + fn columns_read_after_vector_search(&self, also_read: &[String]) -> Vec { + let mut columns: Vec = Vec::new(); + let push = |columns: &mut Vec, name: &str| { + if !columns.iter().any(|existing| existing == name) { + columns.push(name.to_string()); + } + }; + + let projection = &self.projection_plan.physical_projection; + for field in &self.dataset.schema().fields { + if projection.contains_field_id(field.id) { + push(&mut columns, &field.name); + } + } + + if let Some(ordering) = &self.ordering { + for column in ordering { + push(&mut columns, &column.column_name); + } + } + + if let Some(aggregate) = &self.aggregate { + for column in aggregate.required_columns() { + push(&mut columns, &column); + } + } + + for column in also_read { + push(&mut columns, column); + } + + columns + } + + /// The columns an FTS scorer reads out of the plan it is handed, rather than fetching + /// for itself. + /// + /// Two nodes do this, and both sit *between* a vector search and the final + /// projection, so neither is visible to [`Self::columns_read_after_vector_search`]: + /// [`Scanner::fts_rerank`] (`FlatMatchQueryExec`, when FTS is the search and a vector + /// query is the filter) and the refine pass over an FTS query filter + /// (`FlatMatchFilter`, via `FilterPlan::refine_columns`). + /// + /// Which shapes score from their input is decided by [`fts_scores_from_input`], which + /// `FilterPlan::refine_columns` consults for the matching take. Narrowing must not be + /// looser than that rule: declaring a covering column no one reads is precisely the + /// cost this pushdown exists to remove. + /// + /// Names are passed through as the query states them. A dotted FTS field path names + /// no top-level column and so contributes nothing, which under-narrows exactly like a + /// filter on `nested.x` -- the scorer then takes the column from the base table, as it + /// does today. + /// Both callers pass queries `create_plan` has already resolved through + /// [`Self::resolve_full_text_search_query`], so the default unnamed `Match` arrives + /// with its column filled from the FTS-indexed columns -- no fill is needed (or + /// wanted) here. A caller handing this an UNRESOLVED query would silently narrow + /// the filled column away and un-engage covering for the most common query shape, + /// with results staying byte-identical; the unnamed-Match cases in + /// `test_covered_ann_keeps_the_column_an_fts_scorer_reads` pin against that. + fn fts_scored_columns(query: &FullTextSearchQuery) -> Vec { + if !fts_scores_from_input(&query.query) { + return Vec::new(); + } + let mut columns: Vec = query.columns().into_iter().collect(); + // `columns()` is a set; sort so the resolved projection is deterministic across + // runs. The emission order is the index's, but a stable request order keeps plans + // reproducible. + columns.sort_unstable(); + columns + } + + /// `query` with [`Query::covering_projection`] resolved against this scan: the one + /// place a covered index is told what to materialize. + /// + /// Called from the plan sources rather than from `vector_search`, because they still + /// hold the complete [`FilterPlan`]. By the time `vector_search` runs the plan has + /// been split into a prefilter and a refine half and it sees only one of them, so + /// resolving there would forget the filter columns of a postfiltered scan -- which + /// is the default. Deliberately not *also* resolved there as a fallback: two sites + /// computing the same policy from different inputs means the poorer one silently + /// covers for a mistake in the richer one. + /// + /// Always `Some`. The empty case is the point (see + /// [`Self::columns_read_after_vector_search`]); a `vector_search` reached without + /// passing through here keeps `None` and materializes everything, which is the safe + /// direction. + /// + /// `also_read` is for consumers the scan options do not describe -- today the FTS + /// scorers stacked on a vector search (see [`Self::fts_scored_columns`]). Callers with + /// no such consumer pass `&[]`. + /// + /// # The expression filter counts only when it is rechecked above the search + /// + /// A postfiltered scan re-evaluates the whole filter on the search's output: both plan + /// sources call `FilterPlan::make_refine_only`, so `create_plan` takes + /// `FilterPlan::refine_columns` and runs `refine_filter` above the search. Those + /// columns are read from whatever the search emitted, so a covered index should serve + /// them. + /// + /// A *prefiltered* scan is the opposite: the filter is consumed entirely below the + /// search (as a `ScalarIndexQuery` or a filtered read feeding the prefilter), and both + /// sources then call `FilterPlan::disable_refine`, which empties the expression filter + /// so `create_plan`'s `if filter_plan.has_refine()` take never runs. Nothing above the + /// search reads those columns, so naming them here would make the index materialize a + /// payload that the final projection discards. Storage honors this narrowed, + /// physically verified projection, so naming those columns would also cause wasted + /// index reads. + /// + /// The unindexed-fragment and stale-row paths in `knn_combined` do apply the filter, + /// but they read their columns from the base table, not from the search, so they do not + /// change this. + /// + /// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection + fn with_covering_projection( + &self, + query: &Query, + filter_plan: &FilterPlan, + also_read: &[String], + ) -> Query { + let mut query = query.clone(); + let mut columns = if self.prefilter { + Vec::new() + } else { + filter_plan.expr_filter_plan.all_columns() + }; + columns.extend_from_slice(also_read); + // `Some`, unconditionally: an empty list means "materialize no covering column", + // and folding it into `None` restores full materialization with every result + // still correct. See `Query::covering_projection`. + query.covering_projection = + Some(Arc::from(self.columns_read_after_vector_search(&columns))); + query + } + // ANN/KNN search execution node with optional prefilter + // + // `q.covering_projection` is expected to be resolved already (see + // `with_covering_projection`): the ANN node and every flat/fallback path in + // `knn_combined` derive their covering columns from it, and they are unioned against + // a single schema, so a flat path emitting fewer covering columns than the ANN node + // declares fails the projection with a missing column. #[async_recursion] async fn vector_search( &self, @@ -5640,6 +5868,12 @@ impl Scanner { let q = q.clone(); debug_assert!(q.metric_type.is_some()); + // The ANN schema has already been narrowed against every selected segment's + // physical storage. Keep it as the authority for which declared covering columns + // the flat paths must mirror; the vector take below can widen `knn_node` with + // unrelated base-table columns. + let ann_schema = knn_node.schema(); + // Ensure the vector column is present for distance computation. if knn_node.schema().column_with_name(&q.column).is_none() { let vector_projection = self @@ -5649,21 +5883,21 @@ 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(); + // Every flat path below must produce the covering columns the ANN node actually + // serves before it is unioned with `knn_node`. The manifest/query pair supplies the + // declaration order, while `ann_schema` supplies the storage-authoritative subset. + let covering_columns: Vec = effective_covering( + indexed_segments + .first() + .map(|s| s.covering_fields.as_slice()) + .unwrap_or(&[]), + q.covering_projection.as_deref(), + self.dataset.schema(), + )? + .into_iter() + .filter(|field| ann_schema.column_with_name(field.name()).is_some()) + .map(|field| field.name().clone()) + .collect(); let mut columns = vec![q.column.clone()]; if let Some(expr) = filter_plan.full_expr.as_ref() { @@ -6668,7 +6902,8 @@ impl Scanner { prefilter_source, overlay_block, self.external_row_mask.clone(), - )?; + ) + .await?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, options: SortOptions { @@ -6730,7 +6965,8 @@ impl Scanner { prefilter_source.clone(), overlay_block.clone(), self.external_row_mask.clone(), - )?; + ) + .await?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, options: SortOptions { @@ -16629,4 +16865,426 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(i_array.values(), &[expected_i]); } } + + /// Number of rows written before the index build, and again appended after it. + /// + /// Bounded from below on two sides, so this is close to the smallest fixture these + /// measurements can be made on: + /// + /// - PQ refuses to train an 8-bit codebook on fewer than 256 rows. + /// - below ~500 rows the covering column's whole partition fits in one data page, so + /// reading one survivor's covering reads every row's. At 300 rows the streaming + /// branch of [`test_covered_projected_query_serves_covering_from_the_index`] + /// measures **4.51x** its plain arm with the survivors-only gather *working* and + /// **4.52x** with it disabled -- 0.1% apart, and both far past its bound. The signal + /// is not weakened at that size, it is gone. + const COVERED_FIXTURE_ROWS: usize = 500; + + /// PQ codebook training iterations for every covering fixture below. + /// + /// No assertion in these tests depends on codebook quality -- none measures recall, and + /// both arms of every comparison train with the same parameters -- so this is set to + /// the smallest value that trains. + const COVERED_PQ_ITERS: usize = 2; + + /// The rows every covering fixture below is built from. + /// + /// Shared rather than repeated because the byte counts these fixtures produce are only + /// comparable to each other while the rows are identically shaped: `payload` is + /// ~1 KB/row, so one copy of it (~0.5 MB) dominates everything else these queries read + /// (~0.15 MB) and no measurement is swamped by fixed costs. + /// + /// The width is load-bearing in *both* directions. Narrower rows stop dominating; + /// wider ones break the streaming branch, whose covering read is per probed partition + /// rather than per survivor -- at 4 KB/row it measures 2.1x its plain arm with the + /// gather working, over its own bound. Trade rows for width in neither direction. + fn covered_fixture_rows() -> impl arrow_array::RecordBatchReader + Send + 'static { + gen_batch() + .col( + "vector", + array::rand_vec::(Dimension::from(32)), + ) + .col("payload", array::rand_utf8(ByteCount::from(1024), false)) + .col("extra", array::step::()) + .into_reader_rows( + RowCount::from(COVERED_FIXTURE_ROWS as u64), + BatchCount::from(1), + ) + } + + /// A dataset whose `payload` column is wide (~1 KB/row), half of it indexed and half + /// of it not, for the read-byte measurements below. + /// + /// Every property here is load-bearing for + /// [`test_covered_unprojected_query_pays_for_no_covering_copy`]: + /// + /// - The rows are appended *after* the index build, so a vector query unions an ANN + /// branch (which reads the index) with a flat scan of the unindexed fragment + /// (which reads the base table). A covering column is therefore reachable on + /// *two* paths, and scanner-side narrowing removes exactly the flat-scan one. + /// - The two halves are the same size and come from the same generator, so the + /// unindexed half's `payload` cost is a faithful unit for the indexed half's too. + /// - Four IVF partitions rather than one, so the query reaches the real probe path. + async fn covered_scan_fixture(uri: &str, cover_payload: bool) -> Dataset { + let mut dataset = Dataset::write(covered_fixture_rows(), uri, None) + .await + .unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, COVERED_PQ_ITERS); + if cover_payload { + params.covering_columns(vec!["payload".to_string()]); + } + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + dataset.append(covered_fixture_rows(), None).await.unwrap(); + dataset + } + + /// Bytes read while running an ANN query that projects no dataset column at all. + /// + /// The counter is reset first, so the fixture build is excluded and only the query is + /// measured. + async fn unprojected_ann_query_bytes(dataset: &Dataset) -> u64 { + let query = Float32Array::from(vec![0.3f32; 32]); + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project::<&str>(&[]).unwrap(); + let _ = scan.try_into_batch().await.unwrap(); + dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes + } + + /// Bytes read by scanning `payload` over the fragment appended after the index build: + /// exactly the read that scanner-side covering narrowing removes, measured on the same + /// dataset and the same column so it tracks any encoding change the bound must survive. + async fn unindexed_payload_bytes(dataset: &Dataset) -> u64 { + let fragments = dataset.get_fragments(); + assert_eq!( + fragments.len(), + 2, + "fixture must be one indexed fragment plus one appended unindexed fragment" + ); + let unindexed = fragments.last().unwrap(); + assert_eq!( + unindexed.count_rows(None).await.unwrap(), + COVERED_FIXTURE_ROWS, + "the appended fragment must hold the whole unindexed half" + ); + let unindexed = unindexed.metadata().clone(); + + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.with_fragments(vec![unindexed]); + scan.project(&["payload"]).unwrap(); + let _ = scan.try_into_batch().await.unwrap(); + dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes + } + + /// The covering columns a dataset's sole index declares, by field id. + async fn declared_covering_field_ids(dataset: &Dataset) -> Vec { + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 1, "fixture builds exactly one index"); + indices[0].covering_fields.clone() + } + + /// A covered index must make an unrelated query pay for the covering column *not at + /// all*. + /// + /// Without the projection narrowing, a query that reads no covering column still + /// materializes `payload` on both paths it can reach: the index storage serves it for + /// the indexed half, and the flat scan over the unindexed fragment reads it from the + /// base table. The narrowing removes the flat-scan copy; the survivors-only gather + /// removes the index one, so a covered query now reads neither. + /// + /// # The bound is half a copy, because that is what the reachable regression costs + /// + /// This bound was one and a half copies, calibrated when losing the narrowing cost two + /// copies. The survivors-only gather then removed the index-side copy independently, so + /// losing the narrowing costs **one** copy -- under the old bound, which therefore + /// passed with its own subject disabled. Measured on this fixture, each row taken with + /// `covering_projection` forced to `None` / the gather forced onto its whole-range + /// fallback: + /// + /// | | read bytes | copies of `payload` above plain | + /// |---|---|---| + /// | plain index | 140,930 | -- | + /// | covered, narrowed + gathered (today) | 145,034 | 0.008 | + /// | covered, narrowing disabled | 674,978 | 1.03 | + /// | covered, narrowing and gather both disabled | 1,163,967 | 2.01 | + /// + /// Each row's copy count is against its own run's `plain` and `one_copy`; the plain arm + /// itself varies by up to ~10% run to run, which is why the assertion is a ratio. + /// + /// Half a copy sits an order of magnitude above today's cost and half an order below + /// the cheapest regression, so it discriminates without pinning the exact number. + /// A recorded measurement is a claim with an expiry date: any later change that + /// *reduces* what a regression here costs has to re-derive this bound, not just the + /// baseline row. + /// + /// Expressing it in measured copies rather than as a byte constant keeps its margins if + /// unrelated encoding work moves the absolute totals. The subtraction of `plain` is + /// what makes this a covering measurement rather than a query-cost measurement: it + /// removes the vectors, PQ codes and metadata that both indexes read identically. + #[tokio::test] + async fn test_covered_unprojected_query_pays_for_no_covering_copy() { + let plain = covered_scan_fixture("memory://cov_io_one_copy_plain", false).await; + let covered = covered_scan_fixture("memory://cov_io_one_copy_covered", true).await; + + // Guard the fixture itself. An `covering_columns` that silently did nothing would + // make `covered` cost the same as `plain`, and an upper bound on their difference + // would then pass for the wrong reason. + let payload_id = covered.schema().field("payload").unwrap().id; + assert_eq!( + declared_covering_field_ids(&covered).await, + vec![payload_id], + "the covered fixture's index must actually declare `payload` as covering" + ); + assert!( + declared_covering_field_ids(&plain).await.is_empty(), + "the plain fixture's index must declare no covering column" + ); + + let plain_bytes = unprojected_ann_query_bytes(&plain).await; + let covered_bytes = unprojected_ann_query_bytes(&covered).await; + let one_copy = unindexed_payload_bytes(&covered).await; + + let covering_overhead = covered_bytes.saturating_sub(plain_bytes); + let bound = one_copy / 2; + assert!( + covering_overhead < bound, + "a query projecting no covering column must not pay for `payload` at all: \ + covered read {covered_bytes} bytes vs plain {plain_bytes}, an overhead of \ + {covering_overhead} against a bound of {bound} (half the {one_copy}-byte \ + covering copy). Half a copy or more means the scanner stopped narrowing what \ + the covered index materializes." + ); + + // The same two measurements as a plain ratio. Redundant with the bound above + // against today's code -- both fail on a lost narrowing -- but it is the one that + // stays meaningful if `unindexed_payload_bytes` ever stops being a faithful unit + // for the indexed half, and it costs no extra fixture. + assert!( + covered_bytes < plain_bytes * 2, + "a covered index must not materially inflate a query that projects no covering \ + column: covered read {covered_bytes} bytes vs plain {plain_bytes}" + ); + } + + /// The same fixture with every row indexed, so a covering column has exactly one + /// source -- the index -- and a measurement over it is the index's own read rather + /// than a mix of index and base table. + async fn covered_index_only_fixture( + uri: &str, + mut params: VectorIndexParams, + cover_payload: bool, + ) -> Dataset { + let mut dataset = Dataset::write(covered_fixture_rows(), uri, None) + .await + .unwrap(); + if cover_payload { + params.covering_columns(vec!["payload".to_string()]); + } + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset + } + + /// Bytes read while running an ANN query that projects the covering column, with the + /// rows it returned. + /// + /// All four partitions are probed, so the covering read spans several partitions + /// rather than the single one the default `nprobes` would reach. `nprobes` and + /// `query_parallelism` are set after `nearest` because they apply to the query that + /// call installs. + /// + /// `query_parallelism(1)` routes the search through + /// [`VectorIndex::search_partitions`], whose branch is then chosen by the sub-index: + /// the flat sub-index (IVF_PQ) takes the global top-k heap, HNSW has no such heap and + /// takes the streaming send loop. That is how one query shape measures both covering + /// emit paths. + async fn projected_ann_query(dataset: &Dataset) -> (u64, RecordBatch) { + let query = Float32Array::from(vec![0.3f32; 32]); + let _ = dataset.object_store.as_ref().io_stats_incremental(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.nprobes(4); + scan.query_parallelism(1); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let bytes = dataset + .object_store + .as_ref() + .io_stats_incremental() + .read_bytes; + (bytes, batch) + } + + /// The plan for the query [`projected_ann_query`] runs, so the two arms' plans are + /// derived the same way as their byte counts. + async fn projected_ann_plan(dataset: &Dataset) -> String { + let query = Float32Array::from(vec![0.3f32; 32]); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.nprobes(4); + scan.query_parallelism(1); + scan.project(&["payload"]).unwrap(); + scan.explain_plan(true).await.unwrap() + } + + /// The query a covered index exists for: it projects the covering column, and the + /// index serves it without touching the base table. + /// + /// A plain index answers this with a base-table take of the `k` winners, which renders + /// as `LanceRead`; a covered index answers it out of index storage and takes nothing. + /// Both plans are asserted, because a covered plan that quietly fell back to the take + /// would return byte-identical rows -- the covering column is semantically transparent, + /// so results alone cannot tell the two apart. + /// + /// # The two emit paths get different bounds, because they bound the read differently + /// + /// The global top-k heap reads once per partition that owns a *survivor* (at most + /// `min(nprobe, k)` reads); the streaming loop every `IVF_HNSW_*` index takes reads once + /// per *probed* partition, so its covering cost grows with `nprobe` and, at a fixture + /// this size, with the data page a partition's covering lands in. One shared bound would + /// have to be the looser of the two, which would stop discriminating for the heap. Each + /// case therefore carries its own multiple of its own plain arm, each measured against + /// the gather forced onto its whole-range fallback -- the regression the survivors-only + /// gather replaced, and what a pending fragment-reuse remap still falls back to: + /// + /// | branch | plain | covered | ratio | bound | gather disabled | + /// |---|---|---|---|---|---| + /// | global heap (`pq`) | 64,197 | 67,725 | 1.06x | 1.5x | 4.74x | + /// | streaming (`hnsw_pq`) | 95,512 | 147,009 | 1.54x | 3x | 6.42x | + /// + /// Both bounds are order-of-magnitude claims, set so unrelated encoding work moving the + /// absolute totals cannot make them flaky. Neither is a floor on the fixture: shrinking + /// it further collapses the streaming row's ratio into its own regression -- see + /// [`COVERED_FIXTURE_ROWS`]. + /// + /// Higher up the sweep the two separate sharply -- on a 32-partition, 32,768-row index + /// with the same ~1 KB payload, warm, `nprobes = 32`: global heap 15,202 B / 5 requests + /// (flat in `nprobe`), streaming 454,558 B / 106 requests, 2,069 vs 737 us against its + /// own plain arm. That is recorded, not pinned here: this test is the shape check that + /// each branch serves covering from the index at all, on a fixture small enough to + /// stay a unit test. + /// + /// HNSW is built with a small graph for the same reason the PQ codebook trains for two + /// iterations: nothing here measures recall, and both arms use identical parameters. + #[rstest] + #[case::global_heap_branch( + "pq", + VectorIndexParams::ivf_pq(4, 8, 4, MetricType::L2, COVERED_PQ_ITERS), + 3 + )] + #[case::streaming_branch( + "hnsw_pq", + VectorIndexParams::with_ivf_hnsw_pq_params( + MetricType::L2, + IvfBuildParams::new(4), + HnswBuildParams { + max_level: 3, + m: 8, + ef_construction: 30, + ..HnswBuildParams::default() + }, + PQBuildParams::new(4, 8), + ), + 6 + )] + #[tokio::test] + async fn test_covered_projected_query_serves_covering_from_the_index( + #[case] family: &str, + #[case] params: VectorIndexParams, + // Halves: the global heap's bound is 1.5x, the streaming loop's 3x. An integer + // multiple cannot express the former, and 2x would not discriminate for the latter. + #[case] bound_halves: u64, + ) { + let plain = covered_index_only_fixture( + &format!("memory://cov_projected_plain_{family}"), + params.clone(), + false, + ) + .await; + let covered = covered_index_only_fixture( + &format!("memory://cov_projected_covered_{family}"), + params, + true, + ) + .await; + + // The two cases are the two covering emit paths, not two spellings of one: + // `search_partitions` takes the global top-k heap only when the sub-index has + // one, and only IVF_PQ's flat sub-index does. + use lance_index::vector::{flat::index::FlatIndex, hnsw::HNSW, v3::subindex::IvfSubIndex}; + assert!( + FlatIndex::supports_global_topk_heap(), + "the `pq` case must exercise the global-heap branch" + ); + assert!( + !HNSW::supports_global_topk_heap(), + "the `hnsw_pq` case must exercise the streaming branch" + ); + + // Guard the fixture itself: an `covering_columns` that silently did nothing would + // make the covered arm cost what the plain one costs, and the bound below would + // then pass for the wrong reason. + let payload_id = covered.schema().field("payload").unwrap().id; + assert_eq!( + declared_covering_field_ids(&covered).await, + vec![payload_id], + "the covered fixture's index must actually declare `payload` as covering" + ); + + // Each dataset serves exactly one measured query, so both are measured against a + // cold partition cache. A second query on the same dataset would find the + // partitions already resident and read far less. + let (plain_bytes, plain_batch) = projected_ann_query(&plain).await; + let (covered_bytes, covered_batch) = projected_ann_query(&covered).await; + assert_eq!(plain_batch.num_rows(), 10); + assert_eq!(covered_batch.num_rows(), 10); + assert_eq!( + covered_batch + .column_by_name("payload") + .expect("the covered result must carry the projected covering column") + .null_count(), + 0, + "the covering values must come back from the index, not as a null fill" + ); + let bound = plain_bytes * bound_halves / 2; + assert!( + covered_bytes < bound, + "a covered index must serve the column it covers for about what a plain index \ + pays to take it: the {family} branch read {covered_bytes} bytes covered vs \ + {plain_bytes} plain, against a bound of {bound} \ + ({bound_halves} halves of plain). A covering read proportional to the \ + partition's size rather than to its survivors overshoots this by several times." + ); + + let plain_plan = projected_ann_plan(&plain).await; + assert!( + plain_plan.contains("LanceRead"), + "the plain arm must take `payload` from the base table, which is the work the \ + covered arm is expected to avoid; without it the assertion below is vacuous. \ + Plan:\n{plain_plan}" + ); + let covered_plan = projected_ann_plan(&covered).await; + assert!( + !covered_plan.contains("LanceRead"), + "a covered index must serve `payload` from the index, with no base-table take. \ + Plan:\n{covered_plan}" + ); + } } diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index e4eddb6a321..dc3f1ce7bf2 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -346,6 +346,918 @@ async fn test_covered_ann_query_elides_base_table_take() { assert!(batch.column_by_name("extra").is_some()); } +/// A dataset with a covered IVF_PQ index over two covering columns, for the +/// projection-pushdown tests below. +/// +/// Deliberate properties, each load-bearing for at least one assertion: +/// - `covering_columns(["payload", "price"])` is the *reverse* of schema order, so an +/// implementation that walks the schema and filters by membership produces +/// `["price", "payload"]` and fails. +/// - `price` and `payload` have different types and values that are neither the row id +/// nor each other (`price` is negative, `payload` is a string), so a positional read +/// of the storage batch cannot pass for a by-name one. +/// - `extra` is a real column the index does not carry, so "narrowed away" and +/// "never declared" stay distinguishable. +/// - Four well-separated clusters and four partitions, so the runtime probe path is +/// genuinely reached rather than short-circuited (see the module doc above). +async fn covered_two_column_dataset(uri: &str) -> Dataset { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let mut rng = StdRng::seed_from_u64(7); + let mut price = Vec::with_capacity(TOTAL); + 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 index = (cluster * ROWS_PER_CLUSTER + row) as i32; + price.push(-index - 7); + payload.push(format!("p{}", index * 3 + 11)); + extra.push(index * 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("price", DataType::Int32, false), + ArrowField::new("payload", DataType::Utf8, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + vectors, + Arc::new(Int32Array::from(price)), + Arc::new(StringArray::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, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset +} + +/// A query vector near the last cluster's centre of [`covered_two_column_dataset`]. +fn covered_query_vector() -> Float32Array { + let mut values = vec![0.0f32; 16]; + values[0] = 3000.0; + Float32Array::from(values) +} + +/// The ANN search node's declared output schema: `[_distance, _rowid]` plus exactly the +/// covering columns this scan asks the index to materialize. +fn ann_declared_schema(plan: &Arc) -> Vec { + fn find( + plan: &Arc, + ) -> Option> { + if plan.name() == "ANNSubIndexExec" { + return Some(plan.clone()); + } + plan.children().into_iter().find_map(find) + } + let ann = find(plan).expect("the scan must plan an ANN sub-index search"); + ann.schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() +} + +/// The `projection=[...]` list of the (single) flat `LanceScan` in an explained plan. +fn flat_scan_projection(plan: &str) -> Vec { + let line = plan + .lines() + .find(|line| line.contains("LanceScan:")) + .unwrap_or_else(|| panic!("plan has no flat LanceScan node:\n{plan}")); + let start = line + .find("projection=[") + .expect("LanceScan prints a projection") + + "projection=[".len(); + let rest = &line[start..]; + let end = rest.find(']').expect("projection list is bracketed"); + rest[..end] + .split(", ") + .filter(|name| !name.is_empty()) + .map(str::to_string) + .collect() +} + +/// The covering set is per query, not per index: an index carrying `payload` and `price` +/// must be asked for exactly the subset a given scan would otherwise take from the base +/// table, and for *nothing* when the scan reads no dataset column at all. +/// +/// Asserted on the ANN node's declared schema rather than on results, because a covering +/// column is semantically transparent: with the narrowing disabled every case below +/// still returns identical, correct values -- they just arrive via a base-table read. +/// The `Some(vec![])` case (no covering column projected) is the one the whole pushdown +/// exists for, and the one that silently disappears if it is folded into "not computed". +/// +/// This is a plan-shape assertion, decided by a static schema comparison at plan +/// construction, so it needs no clustering and no minimum partition count. +#[tokio::test] +async fn test_covered_ann_declares_only_the_covering_columns_the_query_projects() { + let test_uri = TempStrDir::default(); + let dataset = covered_two_column_dataset(&test_uri).await; + let query = covered_query_vector(); + + let declared_with = + |columns: Vec<&'static str>, filter: Option<&'static str>, prefilter: bool| { + let dataset = dataset.clone(); + let query = query.clone(); + async move { + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&columns).unwrap(); + if let Some(filter) = filter { + scan.prefilter(prefilter); + scan.filter(filter).unwrap(); + } + ann_declared_schema(&scan.create_plan().await.unwrap()) + } + }; + // Postfiltering is the scanner default; the prefilter cases say so explicitly. + let declared = |columns: Vec<&'static str>, filter: Option<&'static str>| { + declared_with(columns, filter, false) + }; + + // Nothing but the search's own output columns: the index declares covering columns, + // but this query needs none of them, so it must ask for none. Collapsing this into + // "no narrowing computed" restores full materialization invisibly. + assert_eq!( + declared(vec![ROW_ID], None).await, + vec![DIST_COL, ROW_ID], + "a query reading no dataset column must ask the index for no covering column" + ); + + // One of two. + assert_eq!( + declared(vec!["price"], None).await, + vec![DIST_COL, ROW_ID, "price"], + "only the covering column this query reads may be materialized" + ); + assert_eq!( + declared(vec!["payload"], None).await, + vec![DIST_COL, ROW_ID, "payload"], + ); + + // Both, in the index's declaration order -- which is the reverse of schema order, so + // this also pins that the emitted order follows the index and not the projection. + assert_eq!( + declared(vec!["price", "payload"], None).await, + vec![DIST_COL, ROW_ID, "payload", "price"], + "covering columns are emitted in the index's declared order" + ); + + // A column the index does not carry narrows to nothing, and is not confused with a + // covering column that was narrowed away. + assert_eq!( + declared(vec!["extra"], None).await, + vec![DIST_COL, ROW_ID], + "an uncovered projection cannot be served by the index" + ); + + // A filter column is read by the refine pass, so it counts even though it is not in + // the output projection. Postfiltering is the default, so this also pins that the + // resolution happens where the whole filter is still visible. + assert_eq!( + declared(vec!["extra"], Some("price < 0")).await, + vec![DIST_COL, ROW_ID, "price"], + "a covering column used only by the filter still saves a base-table read" + ); + + // ...but only when the filter is rechecked above the search. A *prefiltered* scan + // consumes the filter below it and then calls `disable_refine`, so nothing above the + // search reads `price` and materializing it would be pure waste -- free today, a real + // read once storage honours the declaration. The paired invariant test + // (`..._refetches_a_covering_column`, `prefilter_on_covering`) is what proves this + // narrowing does not cost a base-table take; this assertion is what proves it happens. + assert_eq!( + declared_with(vec!["extra"], Some("price < 0"), true).await, + vec![DIST_COL, ROW_ID], + "a prefiltered scan's filter columns are consumed below the search, so the index \ + must not be asked to materialize them" + ); +} + +/// An FTS scorer sitting on top of a vector search reads one more column than the scan +/// options describe: it scores from whatever the search emitted, so the FTS document +/// column is consumed *between* the search and the final projection. Narrowing that is +/// blind to it makes the scorer take the column from the base table -- correct, but +/// exactly the read a covered index exists to avoid. +/// +/// Both directions of the pairing are covered, because they are separate code paths: +/// +/// * FTS is the search and the vector query is the filter -- `Scanner::fts_rerank`; +/// * the vector query is the search and FTS is the (postfiltered) query filter -- +/// the refine pass, via `FilterPlan::refine_columns`. +/// +/// Each is paired with a non-`Match` shape that discriminates rather than merely +/// differing: only `Match` scores from its input, every other shape builds its own plan +/// and joins on `_rowid`. So the `Match` cases must keep `payload` and the others must +/// still narrow it away -- a "fix" that simply stopped narrowing on FTS plans passes the +/// first of each pair and fails the second. +#[tokio::test] +async fn test_covered_ann_keeps_the_column_an_fts_scorer_reads() { + use crate::dataset::scanner::QueryFilter; + + let test_uri = TempStrDir::default(); + let mut dataset = covered_two_column_dataset(&test_uri).await; + dataset + .create_index( + &["payload"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + let vector_query = lance_index::vector::Query { + column: "vector".to_string(), + key: Arc::new(covered_query_vector()), + // Every row passes the vector filter, so the rerank below is exercised on a + // populated input regardless of where the jitter puts any single row. + k: 256, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(MetricType::L2), + use_index: true, + query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + + // Neither scan projects a dataset column, so `payload` can only ever be declared + // because an FTS scorer is known to read it. + + // FTS is the search, the vector query is the filter: `Scanner::fts_rerank`. + let reranked = |fts: FullTextSearchQuery| { + let dataset = dataset.clone(); + let vector_query = vector_query.clone(); + async move { + let mut scan = dataset.scan(); + scan.full_text_search(fts) + .unwrap() + .prefilter(true) + .filter_query(QueryFilter::Vector(vector_query)) + .unwrap(); + scan.project::<&str>(&[]).unwrap(); + scan + } + }; + + // The vector query is the search, FTS is the query filter. Postfiltering is the + // default, and it is the postfilter path whose refine pass scores from the search. + let postfiltered = |fts: FullTextSearchQuery| { + let dataset = dataset.clone(); + async move { + let mut scan = dataset.scan(); + scan.nearest("vector", &covered_query_vector(), 256) + .unwrap(); + scan.minimum_nprobes(4); + scan.filter_query(QueryFilter::Fts(fts)).unwrap(); + scan.project::<&str>(&[]).unwrap(); + scan + } + }; + + // `p11` is the `payload` value of row 0 (see `covered_two_column_dataset`). + let match_query = || { + FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + )) + }; + // Same terms, but a shape that joins on `_rowid` instead of scoring from the input. + let joined_query = || { + FullTextSearchQuery::new_query(FtsQuery::Boolean(BooleanQuery::new(vec![( + Occur::Should, + FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + ), + )]))) + }; + + let scan = reranked(match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the covered document column the FTS rerank scores from must survive the narrowing" + ); + // The plan is not merely well-shaped: it runs and scores. + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the reranked search must return the matching row" + ); + assert!( + batch.column_by_name(SCORE_COL).is_some(), + "an FTS rerank emits a score column" + ); + + let scan = reranked(joined_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID], + "a non-Match rerank joins on _rowid and reads nothing from the search, so the \ + covering column must still be narrowed away" + ); + + let scan = postfiltered(match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the refine pass over an FTS query filter scores from the search too" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the postfiltered search must return the matching row" + ); + + // No non-Match counterpart for this direction: `Scanner::flat_fts_filter` rejects + // every other shape outright, so there is no plan to assert a schema on. The + // Match-only rule is pinned by the rerank pair above, which shares the same resolver + // (`fts_scores_from_input`). + // + // Asserted on the error's *content*, not merely on `is_err()`: this case exists to + // fire when the restriction is lifted, and any unrelated planning failure -- a fixture + // that stopped building the inverted index, a change to `filter_query` validation -- + // would otherwise keep it green while the claim it protects went untested. + let err = postfiltered(joined_query()) + .await + .create_plan() + .await + .expect_err("a non-Match FTS post-filter is not plannable today"); + assert!( + err.to_string().contains("Only Match queries are supported"), + "if a non-Match FTS post-filter ever becomes plannable it needs its own \ + declared-schema case here; got a different error instead: {err}" + ); + + // The DEFAULT shape: a Match that names no column at all. `create_plan` resolves + // both the FTS search query and the FTS query filter through + // `resolve_full_text_search_query`, which fills the column from the FTS-indexed + // columns BEFORE any covering resolution sees the query -- so the filled column + // must reach the covering declaration on both directions. If that resolution were + // ever skipped (or `fts_scored_columns` handed an unresolved query), the most + // common query shape would silently stop engaging covering, taking `payload` from + // the base table on every query with byte-identical results. + let unnamed_match_query = + || FullTextSearchQuery::new_query(FtsQuery::Match(MatchQuery::new("p11".to_string()))); + + let scan = reranked(unnamed_match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "an unnamed Match resolves to the FTS-indexed column, which the covering \ + declaration must include" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the unnamed-Match rerank must return the matching row" + ); + + let scan = postfiltered(unnamed_match_query()).await; + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + "the refine pass fills the unnamed Match's column, so the covering must too" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the unnamed-Match postfilter must return the matching row" + ); +} + +// --------------------------------------------------------------------------------------- +// The covering-elision invariant +// +// `Scanner::columns_read_after_vector_search` is a hand-written enumeration of the columns +// something above the vector search reads. Nothing ties it to the nodes `create_plan` +// actually builds, and under-declaration is silent: results stay byte-identical because +// `Scanner::take` simply fetches the column from the base table instead. Phase 7 shipped +// two readers missing from that enumeration (`FlatMatchQueryExec` and `FlatMatchFilter`), +// both found by reading plan text, not by a failing test. +// +// The oracle is that plan text. For an index that covers a column, a take above the search +// fetching that column *is* the lost elision -- there is nothing else to look for. So the +// tests below plan a matrix of scan shapes over an index covering **every** non-vector +// column and assert, structurally, that no take above the search fetches any of them. +// Any newly inserted reader trips this in whatever shapes the matrix exercises. +// +// This is a test-level guard, not a self-maintaining derivation; a genuinely new plan shape +// still needs its own case. See the task report for the two-pass planning alternative that +// would remove the enumeration outright. +// --------------------------------------------------------------------------------------- + +/// The covering columns that some take *above* the vector search fetches from the base +/// table -- i.e. the elisions this scan shape lost. +/// +/// Returns `None` when the plan contains no ANN search at all, so a caller can tell +/// "nothing was taken" apart from "there was nothing to take above". +/// +/// Only the ancestors of the ANN node are inspected. The prefilter hangs *below* the +/// search (it is a child of `ANNSubIndexExec`) and the flat/unindexed fallback is a +/// sibling branch of the union; both legitimately read from the base table, and neither is +/// on the path this walks. Of the ancestors, only the two nodes that can fetch from the +/// dataset are counted -- a `ProjectionExec` can rename or extract a struct child, which +/// would otherwise look like a newly appeared column. +fn covering_taken_above_search( + plan: &Arc, + covering: &HashSet, +) -> Option> { + if plan.name() == "ANNSubIndexExec" { + return Some(Vec::new()); + } + for child in plan.children() { + let Some(mut taken) = covering_taken_above_search(child, covering) else { + continue; + }; + if matches!(plan.name(), "FilteredReadExec" | "TakeExec") { + let child_schema = child.schema(); + let below: HashSet<&str> = child_schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect(); + for field in plan.schema().fields() { + if covering.contains(field.name()) && !below.contains(field.name().as_str()) { + taken.push(field.name().clone()); + } + } + } + return Some(taken); + } + None +} + +/// The scan shapes the invariant is checked against. Each is a distinct route through +/// `create_plan`'s post-source pipeline, which is where a reader can hide. +#[derive(Debug, Clone, Copy)] +enum CoveredShape { + /// The state the pushdown exists for: nothing read from the dataset at all. + ProjectNothing, + ProjectOneCovering, + ProjectEveryCovering, + /// Postfiltering is the default; the filter is rechecked above the search. + PostfilterOnCovering, + /// Prefiltering consumes the filter below the search, so nothing above reads it. + PrefilterOnCovering, + /// `create_plan` takes the sort keys above the search, twice. + OrderByCovering, + /// Aggregate inputs are taken above the search on their own path. + AggregateOverCovering, + /// Inserts a `take` of the vector column above the ANN node before the flat re-score. + RefineFactor, + LimitOffset, + WithRowAddress, + /// Skips `knn_combined` entirely. + FastSearch, + /// FTS is the search, the vector query is the filter: `FlatMatchQueryExec` scores from + /// the search's own columns. + FtsRerank, + /// The vector query is the search, FTS is the query filter: the refine pass scores + /// from the search's own columns via `FlatMatchFilter`. + FtsPostfilterRefine, +} + +/// A dataset whose covered index carries **every** non-vector column, so that any column +/// any scan shape reads above the search is a covering column and therefore a detectable +/// elision loss. `payload` also carries an inverted index, for the two FTS shapes. +/// +/// Four well-separated clusters and four partitions, so the runtime probe path is genuinely +/// reached rather than short-circuited by a single partition. +async fn fully_covered_dataset(uri: &str, append_unindexed: bool) -> Dataset { + const DIMS: usize = 16; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const TOTAL: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + + let mut rng = StdRng::seed_from_u64(11); + let mut make_batch = |tag: &str| { + let mut price = Vec::with_capacity(TOTAL); + 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 index = (cluster * ROWS_PER_CLUSTER + row) as i32; + price.push(-index - 7); + payload.push(format!("{tag}{}", index * 3 + 11)); + extra.push(index * 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("price", DataType::Int32, false), + ArrowField::new("payload", DataType::Utf8, false), + ArrowField::new("extra", DataType::Int32, false), + ])); + RecordBatch::try_new( + schema, + vec![ + vectors, + Arc::new(Int32Array::from(price)), + Arc::new(StringArray::from(payload)), + Arc::new(Int32Array::from(extra)), + ], + ) + .unwrap() + }; + + let batch = make_batch("p"); + let schema = batch.schema(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(NUM_CLUSTERS, 8, 4, MetricType::L2, 2); + // Every non-vector column. `covering_columns` rejects the indexed vector column itself, + // so this is the widest legal declaration for this schema. + params.covering_columns(vec![ + "payload".to_string(), + "price".to_string(), + "extra".to_string(), + ]); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index( + &["payload"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + if append_unindexed { + // After the build, so `knn_combined` plans its flat fallback branch and the extra + // `take` of the vector column that sits above the ANN node. + let more = make_batch("q"); + let reader = RecordBatchIterator::new(vec![more].into_iter().map(Ok), schema); + dataset.append(reader, None).await.unwrap(); + } + dataset +} + +/// Build the scan for one shape. Kept in one place so the matrix stays a list of names. +async fn plan_covered_shape( + dataset: &Dataset, + shape: CoveredShape, +) -> Arc { + use crate::dataset::scanner::{AggregateExpr, ColumnOrdering, QueryFilter}; + + let query = covered_query_vector(); + let vector_filter = lance_index::vector::Query { + column: "vector".to_string(), + key: Arc::new(query.clone()), + k: 32, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 4, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(MetricType::L2), + use_index: true, + query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: Default::default(), + covering_projection: None, + }; + // `p11` is row 0's payload; any term would do, the shape is what is under test. + let match_query = || { + FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("p11".to_string()).with_column(Some("payload".to_string())), + )) + }; + + let mut scan = dataset.scan(); + // Every shape except the two FTS ones is a plain vector search. + if !matches!( + shape, + CoveredShape::FtsRerank | CoveredShape::FtsPostfilterRefine + ) { + scan.nearest("vector", &query, 10).unwrap(); + scan.minimum_nprobes(4); + } + scan.with_row_id(); + + match shape { + CoveredShape::ProjectNothing => { + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::ProjectOneCovering => { + scan.project(&["payload"]).unwrap(); + } + CoveredShape::ProjectEveryCovering => { + scan.project(&["payload", "price", "extra"]).unwrap(); + } + CoveredShape::PostfilterOnCovering => { + scan.prefilter(false); + scan.filter("price < 0 AND extra >= 0").unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::PrefilterOnCovering => { + scan.prefilter(true); + scan.filter("price < 0 AND extra >= 0").unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::OrderByCovering => { + scan.project::<&str>(&[]).unwrap(); + scan.order_by(Some(vec![ColumnOrdering::asc_nulls_first( + "price".to_string(), + )])) + .unwrap(); + } + CoveredShape::AggregateOverCovering => { + scan.project::<&str>(&[]).unwrap(); + scan.aggregate(AggregateExpr::builder().sum("extra").build()) + .unwrap(); + } + CoveredShape::RefineFactor => { + scan.refine(2); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::LimitOffset => { + scan.limit(Some(5), Some(2)).unwrap(); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::WithRowAddress => { + scan.with_row_address(); + scan.project(&["price"]).unwrap(); + } + CoveredShape::FastSearch => { + scan.fast_search(); + scan.project(&["payload"]).unwrap(); + } + CoveredShape::FtsRerank => { + scan.full_text_search(match_query()) + .unwrap() + .prefilter(true) + .filter_query(QueryFilter::Vector(vector_filter)) + .unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + CoveredShape::FtsPostfilterRefine => { + scan.nearest("vector", &query, 32).unwrap(); + scan.minimum_nprobes(4); + scan.prefilter(false); + scan.filter_query(QueryFilter::Fts(match_query())).unwrap(); + scan.project::<&str>(&[]).unwrap(); + } + } + scan.create_plan().await.unwrap() +} + +/// No take above a covered vector search may fetch a column the index already carries. +/// +/// This is the structural counterpart to +/// `test_covered_ann_declares_only_the_covering_columns_the_query_projects`: that test +/// pins what the *declaration* says for a handful of shapes, this one pins the +/// *consequence* -- that the declaration was complete enough that nothing had to be +/// re-fetched -- across every shape in the matrix, without naming a single column. +/// +/// The assertion is paired, not bare: each case first asserts the shape really planned an +/// ANN search (otherwise "no take above it" is vacuous), then that the set of covering +/// columns taken above it is empty. +/// +/// Run with and without an unindexed fragment, because the fallback rewrites the whole +/// region above the search into a union with an extra take of the vector column. +#[rstest] +#[case::project_nothing(CoveredShape::ProjectNothing)] +#[case::project_one_covering(CoveredShape::ProjectOneCovering)] +#[case::project_every_covering(CoveredShape::ProjectEveryCovering)] +#[case::postfilter_on_covering(CoveredShape::PostfilterOnCovering)] +#[case::prefilter_on_covering(CoveredShape::PrefilterOnCovering)] +#[case::order_by_covering(CoveredShape::OrderByCovering)] +#[case::aggregate_over_covering(CoveredShape::AggregateOverCovering)] +#[case::refine_factor(CoveredShape::RefineFactor)] +#[case::limit_offset(CoveredShape::LimitOffset)] +#[case::with_row_address(CoveredShape::WithRowAddress)] +#[case::fast_search(CoveredShape::FastSearch)] +#[case::fts_rerank(CoveredShape::FtsRerank)] +#[case::fts_postfilter_refine(CoveredShape::FtsPostfilterRefine)] +#[tokio::test] +async fn test_no_take_above_a_covered_search_refetches_a_covering_column( + #[case] shape: CoveredShape, + #[values(false, true)] unindexed_fragment: bool, +) { + let test_uri = TempStrDir::default(); + let dataset = fully_covered_dataset(&test_uri, unindexed_fragment).await; + let covering: HashSet = ["payload", "price", "extra"] + .into_iter() + .map(str::to_string) + .collect(); + + let plan = plan_covered_shape(&dataset, shape).await; + let taken = covering_taken_above_search(&plan, &covering).unwrap_or_else(|| { + panic!( + "{shape:?} (unindexed_fragment={unindexed_fragment}) planned no ANN search, so \ + the invariant below would hold vacuously; either the shape stopped using the \ + index or it does not belong in this matrix. Plan:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ) + }); + assert!( + taken.is_empty(), + "{shape:?} (unindexed_fragment={unindexed_fragment}) re-fetches covering column(s) \ + {taken:?} from the base table above the search. The index already carries them, so \ + some node above the search reads a column \ + `Scanner::columns_read_after_vector_search` does not name. Plan:\n{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); +} + +/// The ANN node and the flat fallback are unioned against a single schema, so they must +/// narrow together. A flat path that emits *fewer* covering columns than the ANN node +/// declares is not a slow plan, it is a failed one -- which is why the pre-pushdown code +/// simply emitted every declared covering column on both sides. +/// +/// Asserts the flat scan's projection positively (exact list), so it fails both when the +/// flat path keeps reading a narrowed-away column and when it stops reading a needed one. +#[tokio::test] +async fn test_covered_ann_flat_fallback_narrows_with_the_index() { + let test_uri = TempStrDir::default(); + let mut dataset = covered_two_column_dataset(&test_uri).await; + + // Append after indexing so `knn_combined` has an unindexed fragment to scan flat. + let schema = Arc::new(ArrowSchema::from(dataset.schema())); + let appended: ArrayRef = Arc::new( + ::try_new_from_values( + Float32Array::from(vec![3000.0f32; 16]), + 16, + ) + .unwrap(), + ); + let append = RecordBatch::try_new( + schema.clone(), + vec![ + appended, + Arc::new(Int32Array::from(vec![-9999])), + Arc::new(StringArray::from(vec!["appended"])), + Arc::new(Int32Array::from(vec![12345])), + ], + ) + .unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(append)], schema.clone()), + None, + ) + .await + .unwrap(); + + let query = covered_query_vector(); + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["price"]).unwrap(); + + // The index really is used -- without this the assertions below would hold just as + // well for a plan that fell back to brute force entirely. + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the scan must use the covered index:\n{plan}" + ); + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "price"], + ); + assert_eq!( + flat_scan_projection(&plan), + vec!["vector", "price"], + "the flat fallback must read the vector it scores plus exactly the covering \ + columns the ANN node declares -- no more (wasted IO) and no fewer (broken \ + union):\n{plan}" + ); + + // And the plan actually runs, with the narrowed covering values arriving by name out + // of a storage batch that still physically carries both covering columns. + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.num_rows() > 0, + "the covered search must return the appended row's cluster" + ); + let prices = batch + .column_by_name("price") + .expect("price is projected") + .as_primitive::(); + assert!( + prices.iter().all(|price| price.unwrap() < 0), + "price values must be the negative fixture values, not row ids or payload \ + positions: {prices:?}" + ); +} + +/// Every covering value must survive a narrowed declaration, by name. The storage batch +/// still carries both covering columns at this point (narrowing storage is a later +/// phase), so a positional read would hand back `payload` where `price` was asked for -- +/// which the fixture's disjoint types and values make impossible to miss. +/// +/// Uses the four-cluster fixture and asserts against `take_rows` ground truth rather +/// than recall: recall cannot discriminate an index read from a brute-force fallback, +/// since a covering column is semantically transparent and a fallback scores 1.0. +#[tokio::test] +async fn test_covered_ann_narrowed_values_match_the_base_table() { + let test_uri = TempStrDir::default(); + let dataset = covered_two_column_dataset(&test_uri).await; + let query = covered_query_vector(); + + let mut scan = dataset.scan(); + scan.nearest("vector", &query, 10).unwrap(); + scan.project(&["payload"]).unwrap(); + scan.with_row_id(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the values below must come from the index, not a fallback scan:\n{plan}" + ); + assert!( + !plan.contains("LanceRead"), + "a covered projection must not take from the base table:\n{plan}" + ); + assert_eq!( + ann_declared_schema(&scan.create_plan().await.unwrap()), + vec![DIST_COL, ROW_ID, "payload"], + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let row_ids = batch + .column_by_name(ROW_ID) + .expect("row ids requested by name") + .as_primitive::() + .values() + .to_vec(); + let from_index = batch + .column_by_name("payload") + .expect("payload is projected") + .as_string::() + .clone(); + + let ground_truth = dataset + .take_rows(&row_ids, dataset.schema().project(&["payload"]).unwrap()) + .await + .unwrap(); + let expected = ground_truth + .column_by_name("payload") + .expect("take returns the requested column") + .as_string::() + .clone(); + assert_eq!( + from_index, expected, + "the narrowed covering column must carry the same values the base table would" + ); +} + /// 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` -- diff --git a/rust/lance/src/dataset/tests/dataset_scanner.rs b/rust/lance/src/dataset/tests/dataset_scanner.rs index 97cecaa6245..22b9e8e69d2 100644 --- a/rust/lance/src/dataset/tests/dataset_scanner.rs +++ b/rust/lance/src/dataset/tests/dataset_scanner.rs @@ -282,6 +282,7 @@ async fn test_vector_filter_fts_search() { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; // Case 1: search with prefilter=true, query_filter=vector([300,300,300,300]) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 616432b6905..967afc3d47d 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -78,6 +78,7 @@ use vector::utils::get_vector_type; mod api; pub(crate) mod append; +pub(crate) mod covering; mod create; pub mod frag_reuse; pub mod mem_wal; diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index adfe0f9e946..9f901c4f360 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -801,36 +801,27 @@ pub async fn merge_indices_with_unindexed_frags<'a>( .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::>()?; + // they survive optimize/merge and partition split/join. Resolution goes + // through `effective_covering` -- the single authority for covering + // resolution -- so it is top-level only and fails loudly on an unresolvable + // id, exactly like the read path: silently rebuilding without the column + // would leave the committed metadata advertising a column the storage + // cannot emit, and a recursive lookup would resolve a (corrupt or foreign) + // nested id to a leaf name the read path errors on. + let covering_columns: Vec = crate::index::covering::effective_covering( + &old_indices[0].covering_fields, + None, + dataset.schema(), + ) + .map_err(|e| { + Error::index(format!( + "Append index: covering declaration of index '{}' cannot be resolved: {e}", + old_indices[0].name + )) + })? + .into_iter() + .map(|field| field.name().clone()) + .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?; diff --git a/rust/lance/src/index/covering.rs b/rust/lance/src/index/covering.rs new file mode 100644 index 00000000000..5f4bb92605b --- /dev/null +++ b/rust/lance/src/index/covering.rs @@ -0,0 +1,283 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Read-path resolution of an index's covering ("included") columns. +//! +//! # Declaration versus capability +//! +//! Covering has two deliberately independent sources of truth: +//! +//! * the **manifest declaration** — `IndexMetadata::covering_fields`, a dependency list of +//! dataset field ids in declaration order; +//! * the **physical capability** — columns the selected segment's storage schema proves it +//! can serve, including their source dataset field ids. +//! +//! The declaration is not evidence that the payload exists. Transitional writers and +//! lifecycle operations may preserve the dependency while omitting or withdrawing the +//! physical values. Readers may skip a base-table take only for the declaration/capability +//! intersection proven across every selected segment. +//! +//! This module resolves only the manifest side. Storage-aware planners must narrow its +//! result through [`VectorIndex::physical_covering_fields`] before adding columns to an +//! executable output schema. +//! +//! Nothing here is vector-specific: `covering_fields` lives on `IndexMetadata`, so the +//! FTS execs resolve covering exactly the same way and must not re-derive it. +//! +//! [`VectorIndex::physical_covering_fields`]: lance_index::vector::VectorIndex::physical_covering_fields + +use arrow_schema::Field as ArrowField; +use lance_core::datatypes::Schema; +use lance_core::{Error, Result}; + +/// The logically requested covering columns for one query: everything the manifest +/// declares, narrowed to that query's covering projection. +/// +/// `projection` is [`Query::covering_projection`], whose three states this function is +/// the primary consumer of: +/// +/// | `projection` | result | +/// |---------------|----------------------------------------------| +/// | `None` | every declared covering column | +/// | `Some(&[])` | empty — the caller must do no covering work | +/// | `Some(cols)` | the declared columns named in `cols` | +/// +/// An empty result is a *contract*, not an incidental outcome: it means the caller skips +/// the covering read altogether rather than projecting an already-loaded batch down to +/// zero columns. Note that `None` and `Some(&[])` therefore differ enormously in cost +/// while being indistinguishable in results. +/// +/// Returns resolved Arrow **fields**, not ids. This is not yet an executable capability: +/// callers that plan a storage read must intersect these fields with the physical segment +/// schemas before declaring them in an output plan. +/// +/// # Order +/// +/// The result follows `covering_fields`, never the order of `projection`. That is the +/// logical output order. Physical storage must prove that it can emit the selected fields +/// in a compatible order; otherwise the reader falls back to the base table. +/// +/// # Errors +/// +/// A declared id absent from the dataset schema's top-level fields is an error, and is +/// reported even when `projection` would have narrowed that column away. The manifest is +/// authoritative for dependencies: such a mismatch means index metadata and schema +/// disagree, and letting a narrowing decision suppress it would turn a corrupt index into +/// one that appears to work for some queries. For the same reason callers must run any cross-segment +/// agreement checks on the *declared* sets before calling this — narrowing first can +/// reduce genuinely disagreeing segments to a subset on which they happen to agree. +/// +/// [`Query::covering_projection`]: lance_index::vector::Query::covering_projection +pub fn effective_covering( + covering_fields: &[i32], + projection: Option<&[String]>, + schema: &Schema, +) -> Result> { + Ok( + effective_covering_with_ids(covering_fields, projection, schema)? + .into_iter() + .map(|(_, field)| field) + .collect(), + ) +} + +/// [`effective_covering`] with each resolved Arrow field still paired to the manifest +/// field id that selected it. Storage-capability checks use this form so logical identity +/// is never reconstructed from a column name. +pub fn effective_covering_with_ids( + covering_fields: &[i32], + projection: Option<&[String]>, + schema: &Schema, +) -> Result> { + if covering_fields.is_empty() { + return Ok(Vec::new()); + } + let mut fields = Vec::with_capacity(covering_fields.len()); + for id in covering_fields { + // Top-level lookup only, deliberately not `Schema::field_by_id`: that recurses + // into struct children, and a nested field matched here would be emitted under + // its leaf name -- a name a covered projection cannot address, and one that can + // collide with an unrelated top-level column. Covering columns are top-level by + // construction; `validate_covering_columns` rejects dotted paths at create time. + let field = schema + .fields + .iter() + .find(|field| field.id == *id) + .ok_or_else(|| { + Error::index(format!( + "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(schema)` is exactly this conversion + // mapped over every field, but it walks the whole table (including nested + // structs) on every covered plan. + if projection.is_none_or(|wanted| wanted.iter().any(|name| name == &field.name)) { + fields.push((*id, ArrowField::from(field))); + } + } + Ok(fields) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; + + /// `price` and `payload` are the covering candidates; `nested.price` exists to prove + /// that a nested field sharing a top-level name is never what gets resolved. + fn schema() -> Schema { + let arrow = ArrowSchema::new(vec![ + ArrowField::new("vec", DataType::Int32, false), + ArrowField::new("price", DataType::Int32, true), + ArrowField::new("payload", DataType::Utf8, true), + ArrowField::new( + "nested", + DataType::Struct(Fields::from(vec![ArrowField::new( + "price", + DataType::Float64, + true, + )])), + true, + ), + ]); + Schema::try_from(&arrow).unwrap() + } + + fn names(fields: &[ArrowField]) -> Vec<&str> { + fields.iter().map(|f| f.name().as_str()).collect() + } + + fn field_id(schema: &Schema, name: &str) -> i32 { + schema.field(name).unwrap().id + } + + #[test] + fn no_projection_keeps_every_declared_column() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let got = effective_covering(&ids, None, &s).unwrap(); + assert_eq!( + names(&got), + vec!["price", "payload"], + "`None` means no narrowing was computed, so every declared column is emitted" + ); + } + + #[test] + fn projection_narrows_to_what_the_query_needs() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let projection = vec!["payload".to_string()]; + let got = effective_covering(&ids, Some(&projection), &s).unwrap(); + assert_eq!(names(&got), vec!["payload"]); + } + + /// The state the whole feature exists for. It is distinct from `None`, which yields + /// the full declared set -- conflating them silently restores full materialization. + #[test] + fn empty_projection_is_not_the_same_as_no_projection() { + let s = schema(); + let ids = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let narrowed = effective_covering(&ids, Some(&[]), &s).unwrap(); + assert!( + narrowed.is_empty(), + "`Some(&[])` means this query needs no covering column at all" + ); + let unnarrowed = effective_covering(&ids, None, &s).unwrap(); + assert_eq!( + names(&unnarrowed), + vec!["price", "payload"], + "`None` must NOT behave like `Some(&[])`; if these two agree the seam is dead" + ); + } + + /// Emission order is the index's declaration order, never the caller's request order + /// and never schema order. Both alternatives pair values with the wrong names. + /// + /// The declared order below is deliberately the reverse of the schema's ascending + /// field-id order (vec=0, price=1, payload=2, ...), and the requested order is + /// deliberately different again. An implementation that walked `schema.fields()` and + /// filtered by membership, or that echoed the request order, produces `["price", + /// "payload"]` and fails. Do not "tidy" either list back into agreement. + #[test] + fn preserves_declaration_order_not_schema_or_request_order() { + let s = schema(); + let declared = vec![field_id(&s, "payload"), field_id(&s, "price")]; + let requested = vec!["price".to_string(), "payload".to_string()]; + let got = effective_covering(&declared, Some(&requested), &s).unwrap(); + assert_eq!( + names(&got), + vec!["payload", "price"], + "declaration order is storage order; any other order mismatches the batch \ + the storage produces" + ); + + // Permuting only the declaration must permute only the result: that is what + // makes this test sensitive to order rather than to membership. + let permuted = vec![field_id(&s, "price"), field_id(&s, "payload")]; + let got = effective_covering(&permuted, Some(&requested), &s).unwrap(); + assert_eq!(names(&got), vec!["price", "payload"]); + } + + #[test] + fn declared_columns_absent_from_the_projection_are_dropped_entirely() { + let s = schema(); + let ids = vec![field_id(&s, "price")]; + let projection = vec!["vec".to_string()]; + assert!( + effective_covering(&ids, Some(&projection), &s) + .unwrap() + .is_empty() + ); + } + + #[test] + fn a_plain_index_declares_nothing_in_every_projection_state() { + let s = schema(); + for projection in [None, Some(&[][..]), Some(&["price".to_string()][..])] { + assert!(effective_covering(&[], projection, &s).unwrap().is_empty()); + } + } + + /// A nested field can carry the same leaf name as a top-level column. Resolving it + /// would emit a `Float64` column named `price` where the covered projection expects + /// the top-level `Int32` one. + #[test] + fn resolves_top_level_fields_only() { + let s = schema(); + let nested_price_id = s + .field("nested") + .unwrap() + .children + .iter() + .find(|f| f.name == "price") + .unwrap() + .id; + let err = effective_covering(&[nested_price_id], None, &s) + .expect_err("a nested field id is not a valid covering declaration"); + assert!(matches!(err, Error::Index { .. })); + + let top_level = effective_covering(&[field_id(&s, "price")], None, &s).unwrap(); + assert_eq!(top_level[0].data_type(), &DataType::Int32); + } + + /// Metadata that disagrees with the schema must fail loudly even when the narrowing + /// would have discarded the offending column: a corrupt index must not appear to work + /// for the subset of queries that happen not to ask for the broken column. + #[test] + fn unresolvable_declared_id_errors_even_when_narrowed_away() { + let s = schema(); + let ids = vec![9999, field_id(&s, "price")]; + for projection in [None, Some(&[][..]), Some(&["price".to_string()][..])] { + let err = effective_covering(&ids, projection, &s) + .expect_err("a declared covering id absent from the schema must be rejected"); + assert!(matches!(err, Error::Index { .. })); + assert!( + err.to_string().contains("9999"), + "error names the offending id, got: {err}" + ); + } + } +} diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index dffb525927d..09e0f99ec5c 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -642,6 +642,42 @@ async fn prepare_vector_segment_build( Ok((element_type, index_type, ivf_params, shuffler)) } +/// Warn when covering is declared on an index type whose search is a pessimisation for it. +/// +/// Covering pays for itself only where the search settles into ONE global top-k heap: there +/// the covering read is bounded by the survivors (see `gather_survivor_covering`). Every HNSW +/// index, and any query with `query_parallelism > 1`, emits per partition instead, so the +/// covering read scales with partitions probed rather than `k` and is issued serially -- +/// measured at 2.8x a plain index's latency warm at nprobe 32 (1.18x cold). That is slower +/// than the base-table take covering exists to elide, so warn at build time rather than let +/// the regression be found in production. Batching the per-partition reads globally would +/// remove this; until then the combination is legal but not recommended. +/// +/// Called after `validate_covering_columns` at every build entry point: warning first would +/// emit this for configurations that are about to be rejected outright, and warning on only +/// one entry point would let a distributed build silently take the shape this exists to flag. +fn warn_if_covering_is_pessimised(index_type: IndexType, column: &str, params: &VectorIndexParams) { + if params.covering_columns.is_empty() + || !matches!( + index_type, + IndexType::IvfHnswFlat | IndexType::IvfHnswPq | IndexType::IvfHnswSq + ) + { + return; + } + log::warn!( + "Index on column '{}' declares covering columns {:?} on a {:?} index. HNSW \ + search emits per partition rather than through one global top-k heap, so \ + covering reads scale with partitions probed rather than k and are issued \ + serially -- about 2.8x a plain index's latency warm at nprobe 32. Prefer \ + IVF_PQ or IVF_FLAT for covering, or expect covered queries on this index to \ + be slower than uncovered ones.", + column, + params.covering_columns, + index_type + ); +} + /// 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 @@ -747,6 +783,7 @@ pub(crate) async fn build_distributed_vector_index( Some(fragment_ids), ) .await?; + warn_if_covering_is_pessimised(index_type, column, params); let stages = ¶ms.stages; let ivf_centroids = ivf_params @@ -1160,6 +1197,7 @@ async fn build_vector_index_impl( let stages = ¶ms.stages; validate_covering_columns(params, dataset, column)?; + warn_if_covering_is_pessimised(index_type, column, params); match index_type { IndexType::IvfFlat => match element_type { @@ -2081,28 +2119,37 @@ 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::>>()?; + // Through `effective_covering` -- the single authority for covering resolution -- and + // not `Schema::field_by_id`: that recurses into struct children while `Field::name` is + // unqualified, so a (corrupt or foreign) nested id would resolve to a bare leaf name, + // and an unrelated top-level column of that name in the target would then be covered + // silently in its place. `append.rs` routes its own rebuild the same way. + let covering_columns: Vec = crate::index::covering::effective_covering( + &source_index.covering_fields, + None, + source_dataset.schema(), + ) + .map_err(|e| { + Error::index(format!( + "covering declaration of index '{}' cannot be resolved against the source \ + dataset schema: {e}", + source_index.name + )) + })? + .into_iter() + .map(|field| field.name().clone()) + .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)?; + warn_if_covering_is_pessimised( + vector_index_type(source_vector_index.as_ref()), + column_name, + ¶ms, + ); let new_uuid = Uuid::new_v4(); let frag_reuse_index = target_dataset @@ -2892,6 +2939,124 @@ mod tests { /// 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. + /// A covering id naming a *nested* field must not resolve to its bare leaf name. + /// + /// `Schema::field_by_id` recurses into struct children and `Field::name` is + /// unqualified, so a nested id resolves to e.g. `price` -- which can collide with an + /// unrelated top-level column in the target and silently build an index covering the + /// wrong column. `effective_covering` is top-level only for exactly this reason, and + /// the optimize path routes through it (`append.rs`); this path must too. + /// + /// Reachable because `validate_covering_columns` rejects dotted paths only at *create* + /// time -- a manifest written by another build is not covered by it. + #[tokio::test] + async fn test_initialize_vector_index_rejects_a_nested_covering_id() { + use arrow_schema::Fields; + + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let target_uri = format!("{}/target", test_dir.as_str()); + + // Source has a struct whose child is named `price`, plus a top-level `price`. + let rows = 300usize; + let nested_price: arrow_array::ArrayRef = + Arc::new(arrow_array::Int64Array::from(vec![1i64; rows])); + let nested: arrow_array::ArrayRef = Arc::new(arrow_array::StructArray::from(vec![( + Arc::new(Field::new("price", ArrowDataType::Int64, true)), + nested_price, + )])); + let top_price: arrow_array::ArrayRef = + Arc::new(arrow_array::Int64Array::from(vec![999i64; rows])); + 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 schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "meta", + ArrowDataType::Struct(Fields::from(vec![Field::new( + "price", + ArrowDataType::Int64, + true, + )])), + true, + ), + Field::new("price", ArrowDataType::Int64, true), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![nested, top_price, vectors]).unwrap(); + let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut source_dataset = Dataset::write(reader, &source_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_pq(4, 8, 2, MetricType::L2, 20); + params.covering_columns(vec!["price".to_string()]); + source_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("vidx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // Forge what another writer could commit: the declaration rebound to the NESTED + // `meta.price` id. Create-time validation never sees this. + let source_dataset = Dataset::open(&source_uri).await.unwrap(); + let nested_id = source_dataset + .schema() + .field("meta.price") + .expect("nested field must exist") + .id; + let mut source_index = source_dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.name == "vidx") + .unwrap() + .clone(); + let top_id = source_dataset.schema().field("price").unwrap().id; + assert_ne!(nested_id, top_id, "the two `price` fields must differ"); + for id in source_index.fields.iter_mut() { + if *id == top_id { + *id = nested_id; + } + } + for id in source_index.covering_fields.iter_mut() { + if *id == top_id { + *id = nested_id; + } + } + + let target_reader = arrow_array::RecordBatchIterator::new( + Vec::>::new(), + schema.clone(), + ); + let mut target_dataset = Dataset::write(target_reader, &target_uri, None) + .await + .unwrap(); + + let err = initialize_vector_index( + &mut target_dataset, + &source_dataset, + &source_index, + &["vector"], + ) + .await + .expect_err("a nested covering id must not resolve to a bare leaf name"); + let msg = err.to_string(); + assert!( + msg.contains(&nested_id.to_string()), + "the error must name the unresolvable id rather than silently covering the \ + unrelated top-level `price`; was: {msg}" + ); + } + #[tokio::test] async fn test_initialize_vector_index_validates_covering_against_target() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index f90d9b89b2b..3067c63aca4 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -43,7 +43,9 @@ 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::{COVERING_FIELD_IDS_KEY, STORAGE_METADATA_KEY}; +use lance_index::vector::storage::{ + COVERING_FIELD_IDS_KEY, PartitionColumns, STORAGE_METADATA_KEY, +}; use lance_index::vector::transform::Flatten; use lance_index::vector::v3::shuffler::{ DEFAULT_PARTITION_WINDOW_BYTES, EmptyReader, IvfShufflerReader, create_ivf_shuffler, @@ -597,8 +599,12 @@ impl IvfIndexBuilder .as_any() .downcast_ref::>() .ok_or(Error::invalid_input("existing index is not IVF index"))?; + // Every column, not just the internal ones: the remapped storage is + // written straight into the new index file, so a covering column + // that is not read here is a covering column the remapped index + // does not have. let part = ivf_index - .load_partition(part_id, false, &NoOpMetricsCollector) + .load_partition_entry_with_covering(part_id) .await?; let storage = part.storage.remap(&mapping)?; @@ -1582,7 +1588,12 @@ impl IvfIndexBuilder // decoded partition per in-flight task while the rest wait on it. let old_data_filter = source.old_data_filter().await?; - let part_storage = existing_index.load_partition_storage(part_id, None).await?; + // Every column, not just the internal ones: these batches are re-written into + // the index being built, so a covering column that is not read here is a + // covering column the new index does not have. + let part_storage = existing_index + .load_partition_storage(part_id, PartitionColumns::All, None) + .await?; let mut part_batches = part_storage.to_batches()?.collect::>(); // for PQ, the PQ codes are transposed, so we need to transpose them back match Q::quantization_type() { @@ -1741,6 +1752,60 @@ impl IvfIndexBuilder .collect() } + /// Refuse to re-stamp a covering payload onto a different logical field. + /// + /// A rebuild (compaction remap, or `optimize_indices` merging existing sources) copies + /// the covering payload out of existing storage **by name**, while the stamp written + /// here is resolved from the **current** dataset schema. If a covered column was + /// rebound between the two -- dropped and re-added under the same name, so the name + /// resolves to a fresh id -- those disagree, and writing the current id over payload + /// built from the old one silently converts a detectable mismatch into an undetectable + /// one. + /// + /// That matters because the read path's protection *is* this stamp: planning compares + /// the stamped source ids against the index's declared `covering_fields` + /// (`common_physical_covering`) and falls back to a base-table take when they differ. + /// Re-stamping makes the comparison succeed, so the stale payload starts being served + /// and the base-table read is elided. Refusing keeps the index in the state the read + /// path already handles correctly. + /// + /// Sources with no covering payload, or whose stamp is absent/malformed (which already + /// makes their payload unservable), are not the case this guards and are skipped. + /// + /// Nor does it fire when this rebuild writes no stamp at all (`about_to_stamp` empty). + /// That is the degraded state -- a declaration cleared by a pre-covering writer while + /// the storage still physically carries the payload -- and it converges to an ordinary + /// index: with no stamp, `physical_covering_fields` yields nothing, the payload is + /// unservable, and reads use the base table. There is no mismatch to launder, and + /// refusing here would make a degraded index permanently un-optimizable + /// (see `test_degraded_covered_index_can_still_be_optimized`). + fn validate_covering_stamp_is_inherited(&self, about_to_stamp: &[i32]) -> Result<()> { + if about_to_stamp.is_empty() { + return Ok(()); + } + for source in &self.existing_indices { + let stamped: Vec = source + .index + .physical_covering_fields()? + .into_iter() + .map(|(source_id, _)| source_id) + .collect(); + if stamped.is_empty() || stamped == about_to_stamp { + continue; + } + return Err(Error::index(format!( + "cannot rebuild this covered index: its existing storage carries covering \ + values stamped with source field ids {stamped:?}, but the covering columns \ + {:?} now resolve to {about_to_stamp:?} in the dataset schema. A covered \ + column was replaced by a different field of the same name, so the stored \ + payload belongs to the old field. Drop and recreate the index rather than \ + rebuilding it.", + self.covering_columns + ))); + } + Ok(()) + } + /// 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 @@ -2066,6 +2131,7 @@ impl IvfIndexBuilder serde_json::to_string(&storage_partition_metadata)?, ); let covering_field_ids = self.covering_field_ids()?; + self.validate_covering_stamp_is_inherited(&covering_field_ids)?; if !covering_field_ids.is_empty() { storage_writer.add_schema_metadata( COVERING_FIELD_IDS_KEY, diff --git a/rust/lance/src/index/vector/fixture_test.rs b/rust/lance/src/index/vector/fixture_test.rs index 2af020e1df2..88c58cc85af 100644 --- a/rust/lance/src/index/vector/fixture_test.rs +++ b/rust/lance/src/index/vector/fixture_test.rs @@ -266,6 +266,7 @@ mod test { query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let idx = make_idx.clone()(expected_query_at_subindex, metric).await; let (partition_ids, _) = idx.find_partitions(&q).unwrap(); diff --git a/rust/lance/src/index/vector/hamming.rs b/rust/lance/src/index/vector/hamming.rs index 3240ab3021b..6a33b235a6f 100644 --- a/rust/lance/src/index/vector/hamming.rs +++ b/rust/lance/src/index/vector/hamming.rs @@ -26,7 +26,7 @@ use lance_index::vector::VectorIndex; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex}; use lance_index::vector::flat::storage::FLAT_COLUMN; use lance_index::vector::ivf::storage::IvfModel; -use lance_index::vector::storage::VectorStore; +use lance_index::vector::storage::{PartitionColumns, VectorStore}; use lance_linalg::distance::{ BinaryHashValues, ClusteringResult, cluster_pairwise_result, extract_binary_hashes_from_fixed_list, pairwise_hamming_distance_binary_parallel, @@ -242,9 +242,10 @@ async fn hamming_clustering_for_ivf_partition_impl( let mut hash_chunks = Vec::new(); let mut num_hashes = 0; for segment in &segments { + // Only the row ids and the binary hashes are read below, both internal columns. let storage = segment .ivf_flat_bin() - .load_partition_storage(partition_id, None) + .load_partition_storage(partition_id, PartitionColumns::Internal, None) .await?; all_row_ids.extend(storage.row_ids().copied()); for batch in storage.to_batches()? { diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index a7745245e70..e274c3b8cad 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -4895,6 +4895,7 @@ mod tests { use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::OldIndexDataFilter; use lance_index::vector::sq::builder::SQBuildParams; + use lance_index::vector::storage::PartitionColumns; use lance_linalg::distance::l2_distance_batch; use lance_testing::datagen::{ generate_random_array, generate_random_array_with_range, generate_random_array_with_seed, @@ -5350,6 +5351,7 @@ mod tests { query_parallelism: lance_index::vector::DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let (partitions, _) = index.find_partitions(&query).unwrap(); let nearest_partition_id = partitions.value(0) as usize; @@ -6878,7 +6880,10 @@ mod tests { ); // PQ code is on residual space - let pq_store = ivf_idx.load_partition_storage(0, None).await.unwrap(); + let pq_store = ivf_idx + .load_partition_storage(0, PartitionColumns::Internal, None) + .await + .unwrap(); pq_store .codebook() .values() diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index fd1a8f79174..00f8c9fdf7b 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -16,13 +16,16 @@ use std::{ }; use crate::index::vector::{ - IndexFileVersion, builder::index_type_string, utils::gather_covering_columns_by_row_id, + IndexFileVersion, + builder::index_type_string, + utils::{gather_covering_columns_by_row_id, row_id_take_indices}, }; use crate::index::{PreFilter, vector::VectorIndex}; use arrow::compute::concat_batches; use arrow_arith::numeric::sub; use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array, cast::AsArray}; -use arrow_schema::DataType; +use arrow_schema::{DataType, Field}; +use arrow_select::take::take_record_batch; use async_trait::async_trait; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::SendableRecordBatchStream; @@ -70,8 +73,11 @@ use lance_index::{ INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME, Index, IndexType, pb, vector::{ DISTANCE_TYPE_KEY, PartitionSearchControl, PreparedPartitionSearchHandle, Query, - VECTOR_RESULT_SCHEMA, ivf::storage::IVF_METADATA_KEY, quantizer::Quantization, - storage::IvfQuantizationStorage, v3::subindex::IvfSubIndex, + VECTOR_RESULT_SCHEMA, + ivf::storage::IVF_METADATA_KEY, + quantizer::Quantization, + storage::{IvfQuantizationStorage, PartitionColumns}, + v3::subindex::IvfSubIndex, }, }; use lance_index::{INDEX_METADATA_SCHEMA_KEY, IndexMetadata}; @@ -167,6 +173,10 @@ struct PreparedPartitionSearch { query: Query, pre_filter: Arc, partition_id: usize, + /// Rows this partition occupies in the storage file. Compared against the loaded + /// partition's row count to decide whether a storage position is also a file + /// position (see [`CoveringGather::WholeRange`]). + partition_rows: usize, partition_centroid: Option, rq_search_cache: Option>, raw_query_context: Option>, @@ -174,6 +184,45 @@ struct PreparedPartitionSearch { _marker: PhantomData<(S, Q)>, } +/// The covering ("included") columns one query needs from one index, resolved once per +/// search. +/// +/// Absent (`None` at the call sites that hold this) for an ordinary index **and** for a +/// query whose covering projection names none of the declared columns: both mean the same +/// thing to the search path -- emit `[_distance, _rowid]` and do no covering read at all. +/// Collapsing the second case into "project an already-read batch down to nothing" is +/// exactly the regression [`Query::covering_projection`] documents. +struct QueryCovering { + /// `[_rowid, ]`, in index declaration order. Both the declared stream + /// schema and the emitted batches are built from this. + schema: arrow_schema::SchemaRef, + /// The same included columns by name and in the same order, for the storage read. + columns: Vec, +} + +/// What the covering gather must read for one partition's search survivors. +enum CoveringGather { + /// The query needs no covering column from this index. + NotNeeded, + /// Read exactly these positions within the partition's row range. Strictly ascending + /// and deduplicated, which is what the reader's take path requires. + Positions(Vec), + /// Positions are not derivable: a deferred fragment-reuse remap drops rows from the + /// loaded partition, so a storage position is no longer the file position it was read + /// from. Read the partition's whole range instead and match by row id. + WholeRange, +} + +/// Where one heap survivor's covering values live, recorded while its partition is still +/// loaded so the gather after the heap settles is a bounded read rather than a re-search. +#[derive(Debug, Clone, Copy)] +struct CoveringLocation { + partition_id: usize, + /// Position within the partition's row range, or `None` when positions are not + /// derivable for that partition (see [`CoveringGather::WholeRange`]). + position: Option, +} + #[derive(Debug)] pub(crate) struct RabitSearchCache { rotated_centroids: Vec, @@ -834,6 +883,7 @@ impl IVFIndex { query: query.clone(), pre_filter, partition_id, + partition_rows: self.storage.partition_size(partition_id), partition_centroid: self.ivf.centroid(partition_id), rq_search_cache: self.rq_search_cache.clone(), raw_query_context, @@ -858,6 +908,7 @@ impl IVFIndex { query: query.clone(), pre_filter, partition_id, + partition_rows: self.storage.partition_size(partition_id), partition_centroid: self.ivf.centroid(partition_id), rq_search_cache: self.rq_search_cache.clone(), raw_query_context, @@ -866,17 +917,23 @@ impl IVFIndex { }) } + /// The CPU half of a prepared partition search: score the partition and locate its + /// survivors' covering rows. Reading those rows is I/O and happens in the caller's + /// async context (see [`IVFIndex::append_covering`]) -- this runs on the CPU pool, + /// where awaiting is not an option. fn run_prepared_partition_search( use_query_residual: bool, use_residual_scratch: bool, prepared: PreparedPartitionSearch, + want_covering: bool, metrics: &dyn MetricsCollector, scratch: &mut QueryScratch, - ) -> Result { + ) -> Result<(RecordBatch, CoveringGather)> { let PreparedPartitionSearch { query, pre_filter, partition_id, + partition_rows, partition_centroid, rq_search_cache, raw_query_context, @@ -913,8 +970,65 @@ impl IVFIndex { residual, scratch, )?; - // Emit covering columns with the per-partition result (batch path; no re-fetch). - Self::append_covering(batch, &part_entry.storage) + let gather = match want_covering { + true => Self::survivor_positions(&batch, &part_entry.storage, partition_rows)?, + false => CoveringGather::NotNeeded, + }; + Ok((batch, gather)) + } + + /// Locate the rows of `batch` (a `[_distance, _rowid]` partition result) within the + /// partition's row range in the storage file. + /// + /// The scan walks the partition's row ids once against the (`<= k`)-sized survivor set + /// with an early exit, rather than building a partition-sized side map per probe. + /// Positions come out ascending and deduplicated because each storage row is visited + /// at most once, which is what the reader's take path requires. + fn survivor_positions( + batch: &RecordBatch, + storage: &Q::Storage, + partition_rows: usize, + ) -> Result { + // A deferred fragment-reuse remap filters rows out of the partition as it is + // loaded, so from that point a storage position is no longer the file position it + // was read from and cannot address the file. + if storage.len() != partition_rows { + return Ok(CoveringGather::WholeRange); + } + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("search result missing row id".to_string()))? + .as_primitive::(); + let mut needed: HashSet = row_ids.values().iter().copied().collect(); + let mut positions = Vec::with_capacity(needed.len()); + Self::locate_survivors(storage, &mut needed, |_, position| positions.push(position)); + Ok(CoveringGather::Positions(positions)) + } + + /// Scan a partition's row ids once against the (`<= k`)-sized `needed` set with an + /// early exit, calling `on_hit(row_id, storage_position)` for each located row. + /// + /// The one shared core of [`Self::survivor_positions`] (per-partition search) and + /// the locator inside `accumulate_prepared_partition_search` (global-heap search). + /// The two callers differ in how they react to a storage/partition misalignment -- + /// early whole-range fallback vs. scan-anyway with `position: None` -- but the scan + /// itself must stay identical: a future change to when positions are valid (e.g. a + /// new remap variant) goes through here, so the two paths cannot silently diverge + /// on the deferred-remap case. Positions come out ascending and deduplicated + /// because each storage row is visited at most once. + fn locate_survivors( + storage: &Q::Storage, + needed: &mut HashSet, + mut on_hit: impl FnMut(u64, u32), + ) { + for (position, row_id) in storage.row_ids().enumerate() { + if needed.remove(row_id) { + on_hit(*row_id, position as u32); + if needed.is_empty() { + break; + } + } + } } #[allow(clippy::too_many_arguments)] @@ -923,7 +1037,8 @@ impl IVFIndex { use_residual_scratch: bool, prepared: PreparedPartitionSearch, heap: &mut BinaryHeap>, - covering_map: &mut HashMap, + covering_locations: &mut HashMap, + want_covering: bool, scratch: &mut QueryScratch, metrics: &dyn MetricsCollector, ) -> Result<()> { @@ -931,6 +1046,7 @@ impl IVFIndex { query, pre_filter, partition_id, + partition_rows, partition_centroid, rq_search_cache, raw_query_context, @@ -970,66 +1086,47 @@ impl IVFIndex { 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. + // Keep covering O(k): record only where the current heap survivors' covering rows + // live, then prune the map to the heap's membership so rows this partition evicted + // drop their entry. Nothing is read here -- the read happens once, after the heap + // has settled, for the survivors that actually made it (see + // `gather_survivor_covering`); buffering every probed partition's covering to serve + // k winners is exactly the O(nprobe * partition_size) cost this phase removes. + // Every heap survivor is located: 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 located 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::(); + if want_covering { 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)) + .filter(|id| !covering_locations.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()?); - } - } + // A deferred fragment-reuse remap filters rows out of the partition as it + // is loaded, so from that point a storage position no longer addresses the + // file and this partition must be re-read as a whole range. + let aligned = part_entry.storage.len() == partition_rows; + Self::locate_survivors(&part_entry.storage, &mut needed, |row_id, position| { + covering_locations.insert( + row_id, + CoveringLocation { + partition_id, + position: aligned.then_some(position), + }, + ); + }); } - 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). + covering_locations.retain(|id, _| heap_ids.contains(id)); + // Invariant: the map locates 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(), + covering_locations.len(), heap_ids.len(), - "covering map must track exactly the heap's survivors (O(k))" + "covering locations must track exactly the heap's survivors (O(k))" ); } Ok(()) @@ -1062,14 +1159,15 @@ impl IVFIndex { 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). + // `[_rowid, ]` covering the heap survivors and nothing else, + // gathered after the heap settled by `gather_survivor_covering`. `None` when the + // gather was not run (no survivor was located). + covering: Option<&RecordBatch>, + // `[_rowid, ]` schema for the covering columns this query needs, + // or None when the index has none or the query needs none. This is a per-query + // constant, so it -- not whether any covering was gathered -- 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(); @@ -1085,38 +1183,41 @@ impl IVFIndex { // 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() { + if row_id_arr.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. + // columns as empty arrays 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(), - )); + columns.push(arrow_array::new_null_array(field.data_type(), 0)); } } 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())?; + // Survivors but nothing gathered for them. Reported rather than filled with + // nulls: a null fill is indistinguishable from genuine nulls once a covering + // column is nullable, so the query would return the right rows with silently + // wrong values. + let covering = covering.ok_or_else(|| { + Error::index(format!( + "index declares covering columns {:?} but none were gathered for the \ + {} result rows", + covering_schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .filter(|name| *name != ROW_ID) + .collect::>(), + row_id_arr.len(), + )) + })?; + // Align the gathered values to the final row ids by row id, never by + // position: the gather returns each partition's rows in file order, which is + // not the heap's order. A survivor the gather did not return is an error + // there, so this cannot silently drop a row. let row_id_u64 = row_id_arr.as_primitive::(); - let included = gather_covering_columns_by_row_id(&combined, row_id_u64)?; + let included = gather_covering_columns_by_row_id(covering, row_id_u64)?; for (field, array) in included { fields.push(field); columns.push(array); @@ -1128,11 +1229,8 @@ 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. - /// The schema search results carry: `[_distance, _rowid]` plus any covering - /// ("included") columns the storage holds, in storage order. Used to declare + /// The schema search results carry: `[_distance, _rowid]` plus the covering + /// ("included") columns of `covering_schema`, 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>, @@ -1150,23 +1248,78 @@ impl IVFIndex { Arc::new(arrow_schema::Schema::new(fields)) } - fn append_covering(batch: RecordBatch, storage: &Q::Storage) -> Result { - let Some(covering) = storage.covering_batch()? else { + /// The covering columns `query` needs from this index, or `None` when there are none + /// to emit. + /// + /// `None` covers two cases that are the same instruction to the search path -- do no + /// covering work at all: an ordinary index, and a covered index whose covering + /// projection this query narrows to nothing. See [`Query::covering_projection`] for + /// why the second must not degrade into "read everything and project it away". + fn query_covering(&self, query: &Query) -> Result> { + let columns = self + .storage + .covering_columns(query.covering_projection.as_deref())?; + if columns.is_empty() { + return Ok(None); + } + Ok(Some(QueryCovering { + schema: self.storage.covering_read_schema(&columns)?, + columns, + })) + } + + /// Widen a `[_distance, _rowid]` partition result with its survivors' covering values, + /// read from the storage file by position. + /// + /// This is the per-partition emit site. It runs in the caller's async context, after + /// the CPU phase has settled which rows survived, so the read is proportional to `k` + /// rather than to the partition size -- and it is deliberately not cached, being a + /// different set of rows for every query. + async fn append_covering( + &self, + partition_id: usize, + batch: RecordBatch, + gather: CoveringGather, + covering: Option<&QueryCovering>, + io_stats: Option, + ) -> Result { + let positions = match &gather { + CoveringGather::NotNeeded => return Ok(batch), + CoveringGather::Positions(positions) => Some(positions.as_slice()), + CoveringGather::WholeRange => None, + }; + let Some(covering) = covering 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); + if row_ids.is_empty() { + // Nothing survived this partition, so there is nothing to read. Still emit the + // covering columns (empty) so every batch on the stream carries the schema the + // exec declared. + for field in covering.schema.fields() { + if field.name() == ROW_ID { + continue; + } + columns.push(arrow_array::new_null_array(field.data_type(), 0)); + fields.push(field.clone()); + } + } else { + let gathered = self + .storage + .take_covering(partition_id, positions, &covering.columns, io_stats) + .await?; + // By row id, never by position: the gather returns the partition's rows in + // file order while the search result is in distance order. A survivor the + // gather did not return is an error there rather than a silent null. + for (field, array) in gather_covering_columns_by_row_id(&gathered, row_ids)? { + fields.push(field); + columns.push(array); + } } Ok(RecordBatch::try_new( Arc::new(arrow_schema::Schema::new(fields)), @@ -1174,6 +1327,70 @@ impl IVFIndex { )?) } + /// Read `[_rowid, ]` for the heap's survivors: one bounded read per + /// partition that still owns one, after the heap has settled. + /// + /// Each partition's covering is narrowed to its own survivors before the next + /// partition is read, so what the loop accumulates is `O(survivors)` and only one + /// partition's covering is ever alive at a time. That matters on the whole-range + /// fallback (a pending fragment-reuse remap puts every partition on it): keeping each + /// partition's whole covering to concatenate at the end would make peak memory + /// `partitions x partition size`, per concurrent query and outside the shared, + /// evictable partition cache -- the opposite of what bounding this read is for. + /// + /// Returns `None` when no survivor was located, which for a covered query means the + /// heap is empty -- [`Self::global_heap_to_batch`] turns "survivors but nothing + /// gathered" into an error rather than a null fill. + async fn gather_survivor_covering( + &self, + locations: &HashMap, + covering: &QueryCovering, + io_stats: Option, + ) -> Result> { + // Group by partition so each contributing partition is read once, not once per + // survivor. `None` positions mark a partition whose positions are not derivable. + let mut by_partition: HashMap>, Vec)> = HashMap::new(); + for (row_id, location) in locations { + let (positions, row_ids) = by_partition + .entry(location.partition_id) + .or_insert_with(|| (Some(Vec::new()), Vec::new())); + row_ids.push(*row_id); + match (positions.as_mut(), location.position) { + (Some(positions), Some(position)) => positions.push(position), + _ => *positions = None, + } + } + let mut batches = Vec::with_capacity(by_partition.len()); + for (partition_id, (mut positions, row_ids)) in by_partition { + if let Some(positions) = positions.as_mut() { + // The reader's take path requires strictly ascending indices; survivors + // arrive in heap order, which is neither sorted nor partition-local. + positions.sort_unstable(); + positions.dedup(); + } + let gathered = self + .storage + .take_covering( + partition_id, + positions.as_deref(), + &covering.columns, + io_stats.clone(), + ) + .await?; + // By row id, never by position: on the fallback `gathered` is the partition's + // whole range in file order, and even on the scattered read the caller's + // survivors are in heap order. A survivor the read did not return is an error + // here rather than a row quietly dropped from the result. + let survivors = UInt64Array::from(row_ids); + let take_idx = row_id_take_indices(&gathered, &survivors)?; + batches.push(take_record_batch(&gathered, &take_idx)?); + } + if batches.is_empty() { + return Ok(None); + } + Ok(Some(concat_batches(&covering.schema, batches.iter())?)) + } + fn preprocess_partition_query( use_query_residual: bool, use_residual_scratch: bool, @@ -1442,8 +1659,15 @@ impl IVFIndex { .get_or_insert_with_key_hit(cache_key, || async { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); - self.load_partition_entry(partition_id, metrics.io_stats()) - .await + // Internal columns only: this entry is cached under a partition id + // alone, so it must not depend on which covering columns the loading + // query wanted. + self.load_partition_entry( + partition_id, + PartitionColumns::Internal, + metrics.io_stats(), + ) + .await }) .await; match &result { @@ -1461,15 +1685,44 @@ impl IVFIndex { info!(target: TRACE_IO_EVENTS, r#type=IO_TYPE_LOAD_VECTOR_PART, index_type="ivf", part_id=partition_id); metrics.record_part_load(); Ok(Arc::new( - self.load_partition_entry(partition_id, metrics.io_stats()) - .await?, + self.load_partition_entry( + partition_id, + PartitionColumns::Internal, + metrics.io_stats(), + ) + .await?, )) } } + /// Load a partition entry carrying every column its storage file holds, bypassing + /// the partition cache in both directions. + /// + /// Rewrite paths need this: the entry they load is written straight back out, so a + /// covering column left unread is a covering column the rewritten index no longer + /// has. The cache is bypassed rather than consulted because its key + /// ([`IVFPartitionKey`]) has no covering component -- a hit would hand back the + /// codes-only entry the search path stores, and an insert would hand a covering-laden + /// entry to the search path. + pub(crate) async fn load_partition_entry_with_covering( + &self, + partition_id: usize, + ) -> Result> { + if partition_id >= self.ivf.num_partitions() { + return Err(Error::index(format!( + "partition id {} is out of range of {} partitions", + partition_id, + self.ivf.num_partitions() + ))); + } + self.load_partition_entry(partition_id, PartitionColumns::All, None) + .await + } + async fn load_partition_entry( &self, partition_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result> { // `concat_batches` indexes the batches by this schema's field positions @@ -1526,16 +1779,21 @@ impl IVFIndex { self.sub_index_metadata[partition_id].clone(), )?; let idx = S::load(batch)?; - let storage = self.load_partition_storage(partition_id, io_stats).await?; + let storage = self + .load_partition_storage(partition_id, columns, io_stats) + .await?; Ok(PartitionEntry::new(idx, storage)) } pub async fn load_partition_storage( &self, partition_id: usize, + columns: PartitionColumns, io_stats: Option, ) -> Result { - self.storage.load_partition(partition_id, io_stats).await + self.storage + .load_partition(partition_id, columns, io_stats) + .await } /// Names of this index's covering ("included") columns, in storage order, or @@ -1732,6 +1990,10 @@ impl VectorIndex for IVFInd self.ivf.num_partitions() } + fn physical_covering_fields(&self) -> Result> { + Ok(self.storage.physical_covering_fields()) + } + #[instrument(level = "debug", skip(self, pre_filter, metrics))] async fn search_in_partition( &self, @@ -1759,7 +2021,10 @@ impl VectorIndex for IVFInd let scratch_pool = self.scratch_pool.clone(); let use_query_residual = self.use_query_residual; let use_residual_scratch = self.use_residual_scratch; - let (batch, local_metrics) = spawn_cpu(move || { + let covering = self.query_covering(&query)?; + let want_covering = covering.is_some(); + let partition_rows = self.storage.partition_size(partition_id); + let (batch, gather, local_metrics) = spawn_cpu(move || { let param = (&query).into(); let refine_factor = query.refine_factor.unwrap_or(1) as usize; let k = query.k * refine_factor; @@ -1786,16 +2051,26 @@ 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)) + // Locate the survivors' covering rows while the partition is still loaded; + // reading them is I/O and cannot happen on the CPU pool. + let gather = match want_covering { + true => Self::survivor_positions(&batch, &part_entry.storage, partition_rows)?, + false => CoveringGather::NotNeeded, + }; + Result::Ok((batch, gather, local_metrics)) }) .await?; local_metrics.dump_into(metrics); - Ok(batch) + self.append_covering( + partition_id, + batch, + gather, + covering.as_ref(), + metrics.io_stats(), + ) + .await } async fn prepare_partition_search( @@ -1812,6 +2087,15 @@ impl VectorIndex for IVFInd )) } + /// The synchronous phase of a prepared partition search. + /// + /// A covered query cannot be served here: since phase 8 the covering values are read + /// from the storage file once the survivors are known, and this entry point exists + /// precisely to run on the CPU pool where that read cannot be awaited. Callers that + /// need covering columns use [`VectorIndex::search_in_partition`] or + /// [`VectorIndex::search_partitions`], both of which own an async context; this + /// reports the mismatch rather than returning rows with their covering columns + /// silently missing. fn search_prepared_partition( &self, prepared: PreparedPartitionSearchHandle, @@ -1820,19 +2104,40 @@ impl VectorIndex for IVFInd let prepared = prepared .downcast::>() .map_err(|_| Error::internal("failed to downcast prepared partition search"))?; - self.scratch_pool.with_scratch(|scratch| { + if let Some(covering) = self.query_covering(&prepared.query)? { + return Err(Error::index(format!( + "search_prepared_partition cannot materialize the covering columns {:?} \ + this query needs: the covering values are read from index storage after \ + scoring, and this is the synchronous phase of a prepared search. Use \ + search_in_partition or search_partitions instead.", + covering.columns, + ))); + } + let (batch, _) = self.scratch_pool.with_scratch(|scratch| { Self::run_prepared_partition_search( self.use_query_residual, self.use_residual_scratch, *prepared, + false, metrics, scratch, ) - }) + })?; + Ok(batch) } + /// False for a covered index, because [`Self::search_prepared_partition`] rejects the + /// queries such an index exists to serve (the covering values are read after scoring, + /// which is I/O, and that method is the synchronous phase). + /// + /// The flag carries no query, so it answers for the index as a whole: a covered index + /// that a particular query narrows to no covering column would in fact be servable + /// there, but advertising `true` and then failing would strand a dispatcher after it + /// had already paid for the partition load. Saying `false` lets it choose + /// [`VectorIndex::search_in_partition`] or [`VectorIndex::search_partitions`] up + /// front. An unreadable covering schema answers `false` for the same reason. fn supports_prepared_partition_search(&self) -> bool { - true + matches!(self.storage.covering_schema(), Ok(None)) } fn auto_query_parallelism(&self, cpu_pool_size: usize) -> usize { @@ -1871,6 +2176,9 @@ impl VectorIndex for IVFInd let prepare_parallelism = get_num_compute_intensive_cpus().max(1); let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; + // Resolved once per search, before `query` is moved into the prepare stream. + let covering = self.query_covering(&query)?; + let want_covering = covering.is_some(); if control.is_none() && S::supports_global_topk_heap() { let heap_capacity = query.k * query.refine_factor.unwrap_or(1) as usize; @@ -1907,35 +2215,52 @@ 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( - use_query_residual, - use_residual_scratch, - prepared, - &mut heap, - &mut covering_map, - scratch, - search_metrics.as_ref(), - ) - .map_err(DataFusionError::from)?; - } - Ok(()) - })?; - Self::global_heap_to_batch(heap, &covering_map, covering_schema.as_deref()) - .map_err(DataFusionError::from) - }) + let (heap, covering_locations) = spawn_cpu( + move || -> DataFusionResult<( + BinaryHeap>, + HashMap, + )> { + let mut heap = BinaryHeap::with_capacity(heap_capacity); + // O(k) side map locating the heap survivors' covering rows, recorded + // in-flight while each partition is loaded (empty for ordinary + // indexes). Nothing is read until the heap has settled. + let mut covering_locations: HashMap = HashMap::new(); + scratch_pool.with_scratch(|scratch| -> DataFusionResult<()> { + for prepared in prepared { + Self::accumulate_prepared_partition_search( + use_query_residual, + use_residual_scratch, + prepared, + &mut heap, + &mut covering_locations, + want_covering, + scratch, + search_metrics.as_ref(), + ) + .map_err(DataFusionError::from)?; + } + Ok(()) + })?; + Ok((heap, covering_locations)) + }, + ) .await?; + // The gather is I/O, so it runs here rather than on the CPU pool: one bounded + // read per contributing partition, for the survivors only. + let gathered = match covering.as_ref() { + Some(covering) => { + self.gather_survivor_covering(&covering_locations, covering, metrics.io_stats()) + .await? + } + None => None, + }; + let batch = Self::global_heap_to_batch( + heap, + gathered.as_ref(), + covering.as_ref().map(|covering| covering.schema.as_ref()), + )?; + // 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(); @@ -1993,6 +2318,16 @@ impl VectorIndex for IVFInd let search_metrics = metrics.clone(); let search_control = control.clone(); let scratch_pool = self.scratch_pool.clone(); + // The covering gather runs in the async half of the search loop, so it needs the + // index and the metrics sink there as well as inside the CPU closure. + let gather_index = self.clone(); + let gather_metrics = metrics.clone(); + // The per-partition batches are widened with the covering columns this query needs + // (`append_covering`), so declare the matching schema -- the global-heap branch + // above already emits its batch's own (covered) schema. Resolved before the search + // loop takes ownership of the covering set. + let result_schema = + Self::covered_result_schema(covering.as_ref().map(|covering| covering.schema.as_ref())); // Search prepared partitions in batches. Each batch is searched in a single // `spawn_cpu` dispatch (amortizing the per-dispatch overhead the single-worker // design in #6475 avoided), but the channel `recv`/`send` stay in async code so @@ -2055,8 +2390,13 @@ impl VectorIndex for IVFInd // cancellable, so abandoning the await leaves the work running.) let cancel_probe = batch_tx.clone(); let search_output = spawn_cpu(move || { - let mut outputs: Vec> = - Vec::with_capacity(prepared_batch.len()); + // Each output carries the partition it came from and where its + // survivors sit in that partition, so the covering read below can + // be a bounded take rather than a whole-partition scan. The read + // itself is I/O and stays out of this CPU closure. + let mut outputs: Vec< + DataFusionResult<(usize, RecordBatch, CoveringGather)>, + > = Vec::with_capacity(prepared_batch.len()); // `stopped` means the whole search should end (an error, an // early-stop signal, or cancellation), not just this batch. let mut stopped = false; @@ -2070,20 +2410,22 @@ impl VectorIndex for IVFInd stopped = true; break; } + let partition_id = prepared.partition_id; match Self::run_prepared_partition_search( use_query_residual, use_residual_scratch, prepared, + want_covering, search_metrics.as_ref(), scratch, ) .map_err(DataFusionError::from) { - Ok(batch) => { + Ok((batch, gather)) => { if let Some(control) = search_control.as_ref() { control.record_batch(&batch); } - outputs.push(Ok(batch)); + outputs.push(Ok((partition_id, batch, gather))); } Err(err) => { outputs.push(Err(err)); @@ -2108,6 +2450,19 @@ impl VectorIndex for IVFInd } }; for output in outputs { + let output = match output { + Ok((partition_id, batch, gather)) => gather_index + .append_covering( + partition_id, + batch, + gather, + covering.as_ref(), + gather_metrics.io_stats(), + ) + .await + .map_err(DataFusionError::from), + Err(err) => Err(err), + }; if batch_tx.send(output).await.is_err() { return; } @@ -2127,10 +2482,6 @@ 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( result_schema, ReceiverStream::new(batch_rx), @@ -2326,7 +2677,7 @@ mod tests { storage::{RABIT_BLOCKED_EX_CODE_COLUMN, RabitQuantizationMetadata, RabitQueryEstimator}, transform::{EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN}, }; - use lance_index::vector::storage::VectorStore; + use lance_index::vector::storage::{PartitionColumns, VectorStore}; use lance_index::vector::v3::subindex::IvfSubIndex; use crate::dataset::{ @@ -2335,7 +2686,8 @@ mod tests { use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; use crate::index::vector::ivf::v2::{ - IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, PartitionEntry, + CoveringLocation, 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; @@ -3213,7 +3565,11 @@ mod tests { .unwrap(); let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; - let storage = ctx.ivf().load_partition_storage(0, None).await.unwrap(); + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); assert!( storage.batch().column_by_name("id").is_some(), "included column 'id' should be co-located in IVF_PQ partition storage" @@ -3282,7 +3638,11 @@ mod tests { 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(); + let storage = ctx + .ivf() + .load_partition_storage(p, PartitionColumns::All, None) + .await + .unwrap(); total_stored += storage.batch().num_rows(); has_id |= storage.batch().column_by_name("id").is_some(); } @@ -4374,7 +4734,11 @@ mod tests { // 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(); + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); assert!( storage.batch().column_by_name("id").is_none(), "merged storage must drop payload the metadata no longer declares" @@ -4767,24 +5131,17 @@ mod tests { ); } - /// 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. + /// A covering declaration is not proof that every segment already stores the payload. + /// This models a transitional distributed build: both segments carry the same logical + /// declaration, but the new segment was produced by a writer that emitted only the + /// vector-index storage. Commit must preserve that declaration, and planning must see + /// the missing physical capability and fetch `id` from the base table for all results. #[tokio::test] - async fn test_commit_existing_index_segments_rejects_covering_fields_disagreement() { + async fn test_declared_but_physically_absent_covering_field_falls_back_to_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 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()]); @@ -4799,6 +5156,11 @@ mod tests { .await .unwrap(); + let original_indices = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(original_indices.len(), 1); + let declared_fields = original_indices[0].fields.clone(); + let declared_covering_fields = original_indices[0].covering_fields.clone(); + let original_fragment_ids: HashSet = dataset .get_fragments() .iter() @@ -4816,7 +5178,7 @@ mod tests { // 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 + let mut transitional_segment = dataset .create_index_builder(&["vector"], IndexType::Vector, &uncovered_params) .name(INDEX_NAME.to_string()) .fragments(new_fragment_ids) @@ -4825,22 +5187,110 @@ mod tests { .await .unwrap(); assert!( - uncovered_segment.covering_fields.is_empty(), + transitional_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]) + // The logical declaration is intentionally independent of the payload written by + // this segment. Do not rewrite its physical auxiliary storage. + transitional_segment.fields = declared_fields; + transitional_segment.covering_fields = declared_covering_fields; + dataset + .commit_existing_index_segments(INDEX_NAME, "vector", vec![transitional_segment]) .await - .unwrap_err(); + .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"), + "one segment cannot physically serve the declared 'id', so the query must use a \ + base-table take; plan was:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + let row_ids = batch[ROW_ID].as_primitive::(); + let ids = batch["id"].as_primitive::(); + let projection = crate::dataset::ProjectionRequest::from_columns(["id"], dataset.schema()); + let truth = dataset + .take_rows(row_ids.values(), projection) + .await + .unwrap(); + let truth_ids = truth["id"].as_primitive::(); + assert_eq!( + ids, truth_ids, + "fallback values must come from the base table" + ); + } + + /// A declaration may be wider than the physical payload without disabling the part + /// storage can prove. The index carries `id`; `tag` is added to the declaration only. + /// Planning must keep `id` covered and take exactly `tag` from the base table. + #[tokio::test] + async fn test_physical_covering_subset_serves_only_proven_columns() { + 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; + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![( + "tag".to_string(), + "CAST(id AS BIGINT) + 1000".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let id_field_id = dataset.schema().field("id").unwrap().id; + let tag_field_id = dataset.schema().field("tag").unwrap().id; + let mut params = VectorIndexParams::ivf_flat(4, DistanceType::L2); + params.covering_columns(vec!["id".to_string()]); + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name(INDEX_NAME.to_string()) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(segment.covering_fields, vec![id_field_id]); + + // Widen only the logical dependency. The auxiliary storage remains physically + // capable of serving `id` and has no `tag` column. + segment.fields.push(tag_field_id); + segment.covering_fields.push(tag_field_id); + dataset + .commit_existing_index_segments(INDEX_NAME, "vector", vec![segment]) + .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.project(&["id", "tag"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); assert!( - err.to_string() - .contains("a logical index cannot mix declarations"), - "unexpected error: {err}" + plan.contains("LanceRead") && plan.contains("projection=[tag]"), + "only the physically absent 'tag' should use a base-table take; plan was:\n{plan}" ); + + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch["id"].as_primitive::(); + let tags = batch["tag"].as_primitive::(); + for (id, tag) in ids.values().iter().zip(tags.values()) { + assert_eq!(*tag, *id as i64 + 1000); + } } /// Same payoff but forcing the BATCH/late-search path: `k` larger than one @@ -5184,7 +5634,11 @@ mod tests { // 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(); + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); assert!( storage.batch().column_by_name("id").is_some(), "retrained storage should re-materialize covering column 'id'" @@ -5255,6 +5709,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let partitions = Arc::new(UInt32Array::from(vec![0u32, 1, 2, 3])); let dists = Arc::new(Float32Array::from(vec![0.0f32; 4])); @@ -5444,13 +5899,10 @@ mod tests { 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(); + // Empty heap and nothing gathered, but the query DOES want covering columns. + let batch = + IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), None, Some(&covering)) + .unwrap(); assert_eq!(batch.num_rows(), 0); assert_eq!( batch.num_columns(), @@ -5461,8 +5913,7 @@ mod tests { // Ordinary (non-covered) index: bare `[_distance, _rowid]`. let plain = - IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), &HashMap::new(), None) - .unwrap(); + IvfPq::global_heap_to_batch(std::collections::BinaryHeap::new(), None, None).unwrap(); assert_eq!(plain.num_columns(), 2); } @@ -5533,7 +5984,11 @@ mod tests { // 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(); + let storage = ctx + .ivf() + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); assert!( storage.batch().column_by_name("id").is_some(), "remapped storage should keep covering column 'id'" @@ -5709,49 +6164,1481 @@ mod tests { } } - /// 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. + /// A current writer stores its physical covering columns in declaration order. The + /// reader now verifies this independently and falls back if a transitional segment + /// differs, but a newly built index should retain the optimization for every declared + /// column rather than relying on that safety net. /// - /// 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() - ))] + /// Declaration order here is `[payload, price]`: the reverse of both the dataset's + /// field-id order (`id`, `price`, `payload`, `vector`) and the projection order the + /// query asks for. An implementation that sorted by field id, walked the schema, or + /// echoed the request order would produce `[price, payload]` and fail. + /// + /// The two covering columns carry different Arrow types and values that are + /// deliberately disjoint from the row ids and from each other (`price` is negative, + /// `payload` is a string): a fixture where a covering value equals its row id makes a + /// positional-for-by-name substitution invisible. #[tokio::test] - async fn test_covered_survives_deferred_remap_frag_reuse( - #[case] mut params: VectorIndexParams, - ) { + async fn test_covering_declaration_and_storage_agree_on_order() { + use crate::index::covering::effective_covering; + use arrow_array::{Int32Array, StringArray, types::Int32Type}; + use lance_index::vector::storage::VectorStore; + const INDEX_NAME: &str = "vector_idx"; + 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(); - 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(); + + // Well-separated clusters, one per partition, so the ANN path is genuinely + // exercised rather than degenerating into a single-partition scan. + let mut flat = Vec::with_capacity(TOTAL * DIMS); + for row in 0..TOTAL { + let center = (row / ROWS_PER_CLUSTER) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % ROWS_PER_CLUSTER) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let ids = Arc::new(UInt64Array::from_iter_values(0..TOTAL as u64)); + let prices = Arc::new(Int32Array::from_iter_values( + (0..TOTAL as i32).map(|i| -i - 7), + )); + let payloads = Arc::new(StringArray::from_iter_values( + (0..TOTAL).map(|i| format!("p{}", i * 3 + 11)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("price", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![ids, prices, payloads, vectors.clone()]) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(NUM_CLUSTERS, DistanceType::L2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // 1. The manifest records the declaration in the order it was requested. + let price_id = dataset.schema().field("price").unwrap().id; + let payload_id = dataset.schema().field("payload").unwrap().id; + assert!( + price_id < payload_id, + "the fixture needs declaration order to differ from field-id order" + ); + let index_meta = dataset + .load_indices_by_name(INDEX_NAME) + .await + .unwrap() + .into_iter() + .next() + .expect("index metadata"); + assert_eq!(index_meta.covering_fields, vec![payload_id, price_id]); + + // 2. The read path's manifest-side resolution follows that order, whichever order + // the query requests its columns in. + let requested = ["price".to_string(), "payload".to_string()]; + let declared = effective_covering( + &index_meta.covering_fields, + Some(&requested), + dataset.schema(), + ) + .unwrap(); + let declared_names: Vec<&str> = declared.iter().map(|f| f.name().as_str()).collect(); + assert_eq!(declared_names, vec!["payload", "price"]); + + // 3. The storage-side resolution -- `covering_field_indices`, derived from each + // storage's `INTERNAL_COLUMNS` -- lists the same columns in the same order. + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let storage = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex") + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap(); + let storage_schema = storage.schema().clone(); + let storage_names: Vec<&str> = storage + .covering_field_indices() + .into_iter() + .map(|i| storage_schema.field(i).name().as_str()) + .collect(); + assert_eq!( + storage_names, declared_names, + "storage order must equal declaration order; if these diverge the covered \ + values are emitted under the wrong names" + ); + + // 4. End to end: the covered values must match the base table, matched BY NAME. + let q = vectors.value(0); + let q = q.as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vector", q, 10).unwrap(); + scan.nprobes(NUM_CLUSTERS); + scan.with_row_id(); + scan.project(&["price", "payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "the covered query must go through the index; plan was:\n{plan}" + ); + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + + let row_ids: Vec = batch + .column_by_name(ROW_ID) + .expect("row id column") + .as_primitive::() + .values() + .to_vec(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["price", "payload"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let got_prices = batch + .column_by_name("price") + .unwrap() + .as_primitive::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + let got_payloads = batch.column_by_name("payload").unwrap().as_string::(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + for (i, row_id) in row_ids.iter().enumerate() { + assert_eq!( + got_prices.value(i), + truth_prices.value(i), + "row {i} (row_id {row_id}): covered price landed under the wrong name" + ); + assert_eq!( + got_payloads.value(i), + truth_payloads.value(i), + "row {i} (row_id {row_id}): covered payload landed under the wrong name" + ); + } + } + + /// Counts the cache hit/miss signal `IVFIndex::load_partition` reports, which is how + /// the tests below observe whether two searches shared one partition entry. + #[derive(Default)] + struct CacheCountingMetrics { + hits: AtomicUsize, + misses: AtomicUsize, + } + + impl lance_index::metrics::MetricsCollector for CacheCountingMetrics { + fn record_parts_loaded(&self, _num_parts: usize) {} + fn record_index_loads(&self, _num_indexes: usize) {} + fn record_comparisons(&self, _num_comparisons: usize) {} + fn record_index_cache_hits(&self, num_hits: usize) { + self.hits.fetch_add(num_hits, Ordering::Relaxed); + } + fn record_index_cache_misses(&self, num_misses: usize) { + self.misses.fetch_add(num_misses, Ordering::Relaxed); + } + } + + /// A covered IVF_FLAT fixture with one well-separated cluster per partition and two + /// covering columns of different Arrow types. The covering values are deliberately + /// disjoint from the row ids (`price` is negative, `payload` is a string) -- a fixture + /// where a covering value equals its row id hides a positional-for-by-name + /// substitution. Row id equals row offset, so the caller can compute ground truth + /// against `vectors` directly. + async fn covered_flat_fixture( + uri: &str, + index_name: &str, + num_clusters: usize, + rows_per_cluster: usize, + ) -> (Dataset, Arc) { + use arrow_array::{Int32Array, StringArray}; + + const DIMS: usize = 16; + let total = num_clusters * rows_per_cluster; + + let mut flat = Vec::with_capacity(total * DIMS); + for row in 0..total { + let center = (row / rows_per_cluster) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % rows_per_cluster) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let prices = Arc::new(Int32Array::from_iter_values( + (0..total as i32).map(|i| -i - 7), + )); + let payloads = Arc::new(StringArray::from_iter_values( + (0..total).map(|i| format!("p{}", i * 3 + 11)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![prices, payloads, vectors.clone()]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(num_clusters, DistanceType::L2); + params.covering_columns(vec!["payload".to_string(), "price".to_string()]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(index_name.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + (dataset, vectors) + } + + /// A partition loaded for search carries only the storage's own columns, so the entry + /// it is cached under -- `IVFPartitionKey { partition_id }`, which has no covering + /// component -- is the same entry for every query. That is what makes the key correct + /// rather than accidentally correct: the naive narrowing (read whichever covering + /// columns *this* query asked for) would let the first query to touch a partition + /// decide what every later query finds there. + /// + /// Two searches asking for disjoint covering subsets run against the same index, in + /// order. The second must reuse the entry the first populated, that shared entry must + /// carry no covering column at all -- an entry holding `payload` but not `price` is + /// exactly the unsound state ruled out here -- and both must return the partition's + /// true nearest neighbours. + /// + /// The line under test is `PartitionColumns::Internal` in `load_partition_entry`. + /// Switching it to `PartitionColumns::All` makes the covering assertions below fail. + #[tokio::test] + async fn test_partition_cache_entry_is_independent_of_the_query_covering_subset() { + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const K: usize = 10; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query_key = vectors.value(0); + let query_for = |covering: &[&str]| Query { + column: "vector".to_string(), + key: query_key.clone(), + k: K, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + 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(), + covering_projection: Some(covering.iter().map(|c| c.to_string()).collect()), + }; + + let payload_metrics = CacheCountingMetrics::default(); + let payload_result = index + .search_in_partition( + 0, + &query_for(&["payload"]), + Arc::new(NoFilter), + &payload_metrics, + ) + .await + .unwrap(); + assert_eq!( + ( + payload_metrics.misses.load(Ordering::Relaxed), + payload_metrics.hits.load(Ordering::Relaxed) + ), + (1, 0), + "the first search must populate the entry rather than find one already there" + ); + + let price_metrics = CacheCountingMetrics::default(); + let price_result = index + .search_in_partition( + 0, + &query_for(&["price"]), + Arc::new(NoFilter), + &price_metrics, + ) + .await + .unwrap(); + assert_eq!( + ( + price_metrics.hits.load(Ordering::Relaxed), + price_metrics.misses.load(Ordering::Relaxed) + ), + (1, 0), + "the second search must reuse the first one's entry -- if it loaded its own, \ + the entries are query-specific and nothing below proves the key is sound" + ); + + // The one entry the two searches shared holds no covering column at all, so there + // is no subset for one query to have fixed on another's behalf. + let entry = index + .load_partition(0, true, &NoOpMetricsCollector) + .await + .unwrap(); + let part = entry.as_ref(); + assert!( + part.storage.covering_field_indices().is_empty(), + "a cached partition must carry no covering column; it carries {:?}", + part.storage.schema() + ); + assert!( + part.storage.covering_batch().unwrap().is_none(), + "covering_batch() on a codes-only partition must report none" + ); + + // Both searches answer correctly over that shared entry. Ground truth is the + // partition's own rows ranked by distance to the query, computed from the base + // table rather than from the index. + let partition_row_ids: Vec = index + .load_partition_storage(0, PartitionColumns::All, None) + .await + .unwrap() + .row_ids() + .copied() + .collect(); + assert!( + partition_row_ids.len() > K, + "partition 0 must hold more rows than k, or the ranking below is vacuous" + ); + let query_values = query_key.as_primitive::().values().to_vec(); + let mut truth: Vec<(u64, f32)> = partition_row_ids + .iter() + .map(|row_id| { + let vector = vectors.value(*row_id as usize); + let distance = vector + .as_primitive::() + .values() + .iter() + .zip(query_values.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum::(); + (*row_id, distance) + }) + .collect(); + truth.sort_by(|a, b| a.1.total_cmp(&b.1)); + let expected: Vec = truth.iter().take(K).map(|(row_id, _)| *row_id).collect(); + + for (label, result) in [("payload", &payload_result), ("price", &price_result)] { + let mut got: Vec = result + .column_by_name(ROW_ID) + .expect("search result carries row ids") + .as_primitive::() + .values() + .to_vec(); + got.sort_unstable(); + let mut want = expected.clone(); + want.sort_unstable(); + assert_eq!( + got, want, + "the {label} query's neighbours must be the partition's true nearest k" + ); + assert!(result.column_by_name(DIST_COL).is_some()); + } + + // The covering half of "both get right answers": each query receives exactly the + // column it projected, gathered for its survivors out of a shared entry that holds + // neither. `price` is `-row_id - 7` and `payload` is `p{row_id * 3 + 11}`, so a + // value taken from the wrong row -- or from the wrong column -- cannot coincide. + for (label, result, other) in [ + ("payload", &payload_result, "price"), + ("price", &price_result, "payload"), + ] { + assert!( + result.column_by_name(other).is_none(), + "the {label} query must not materialize {other}: a covering column the \ + query never reads is the cost this narrowing exists to avoid, and it is \ + invisible in results" + ); + let row_ids = result + .column_by_name(ROW_ID) + .expect("search result carries row ids") + .as_primitive::(); + let covering = result + .column_by_name(label) + .unwrap_or_else(|| panic!("the {label} query must materialize {label}")); + for i in 0..row_ids.len() { + let row_id = row_ids.value(i); + match label { + "price" => assert_eq!( + covering + .as_primitive::() + .value(i), + -(row_id as i32) - 7, + "row {i} (row_id {row_id}): gathered price belongs to another row" + ), + _ => assert_eq!( + covering.as_string::().value(i), + format!("p{}", row_id * 3 + 11), + "row {i} (row_id {row_id}): gathered payload belongs to another row" + ), + } + } + } + } + + /// A covered IVF_FLAT fixture built for the survivor gather: one well-separated cluster + /// per partition, a **nullable** covering column, and covering values disjoint from the + /// row ids. + /// + /// Both of those properties are load-bearing and no other covered fixture in this suite + /// has them. Every other one declares its covering columns non-nullable, which is why a + /// covered query that returned a silent NULL for every survivor stayed invisible until + /// it was probed deliberately -- non-nullable columns turn it into a loud arrow error + /// instead. And a fixture whose covering value equals its row id cannot tell a by-name + /// gather from a positional one. + /// + /// `payload` is deliberately wide (~1 KB/row). The gather exists for the regime where + /// the covering payload dwarfs the quantization code; at a narrow width the scattered + /// and sequential reads cost almost the same and nothing can be measured about them. + /// + /// Returns the dataset and its vectors; row id equals row offset, so ground truth can + /// be computed against `vectors` directly. + async fn covered_gather_fixture( + uri: &str, + index_name: &str, + num_clusters: usize, + rows_per_cluster: usize, + ) -> (Dataset, Arc) { + use arrow_array::{Int32Array, StringArray}; + + const DIMS: usize = 16; + let total = num_clusters * rows_per_cluster; + + let mut flat = Vec::with_capacity(total * DIMS); + for row in 0..total { + let center = (row / rows_per_cluster) as f32 * 50.0; + for d in 0..DIMS { + flat.push(center + (row % rows_per_cluster) as f32 * 0.001 + d as f32 * 0.0001); + } + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + // Decreasing, and offset by 1000, so `tag` is never equal to its row id. + let tags = Arc::new(Int32Array::from_iter_values( + (0..total as i32).map(|i| 1000 - i), + )); + // Every third row is NULL. With k = 10 the result always contains both a NULL and + // a non-NULL note, so neither case is vacuous. + let notes = Arc::new(StringArray::from_iter((0..total).map(|i| match i % 3 { + 0 => None, + _ => Some(format!("n{}", i * 7 + 3)), + }))); + let payloads = Arc::new(StringArray::from_iter_values( + (0..total).map(|i| format!("{i:04}").repeat(256)), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("tag", DataType::Int32, false), + Field::new("note", DataType::Utf8, true), + Field::new("payload", DataType::Utf8, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![tags, notes, payloads, vectors.clone()]) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, uri, None).await.unwrap(); + + let mut params = VectorIndexParams::ivf_flat(num_clusters, DistanceType::L2); + params.covering_columns(vec![ + "note".to_string(), + "tag".to_string(), + "payload".to_string(), + ]); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(index_name.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + (dataset, vectors) + } + + /// The gather reads the survivors' rows and nothing else -- until they stop being a + /// small fraction of the partition, at which point `Indices` degenerates into a + /// scattered read of nearly everything and the sequential read it replaced is cheaper. + /// + /// Both branches are exercised here against the same real index, and the assertion + /// distinguishes them directly rather than by cost: the scattered read returns exactly + /// the rows asked for, the sequential fallback returns the partition's whole range + /// (the caller aligns by row id either way, so both are correct). A test that passed on + /// either branch would prove nothing about the threshold, so the two row counts are + /// asserted separately and the byte counts are asserted on top. + /// + /// The line under test is `covering_read_for`'s comparison against + /// `COVERING_SCATTERED_READ_MAX_PERCENT`. Neutralise the *comparison*, not the constant: + /// `few` and `many` are derived from the same constant, so moving it drags both sides of + /// the threshold with it and the behavioural assertions stay satisfied. Collapsing + /// `covering_read_for` to `CoveringRead::Sequential` fails this test, and so does + /// collapsing it to `CoveringRead::Scattered` for every non-empty partition. + #[tokio::test] + async fn test_covering_gather_reads_only_the_survivors_rows_until_the_threshold() { + use lance_index::vector::storage::COVERING_SCATTERED_READ_MAX_PERCENT; + use lance_io::scheduler::IoStats; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, _) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + // The largest partition, not partition 0: IVF centroid initialisation is unseeded, + // so any individual partition can come out empty and `partition_size(0) == 0` makes + // both sides of the threshold the same set. Over `NUM_CLUSTERS` partitions the + // largest necessarily holds at least `ROWS_PER_CLUSTER`, which keeps the guard below + // satisfied for every assignment rather than for the lucky ones. + let partition_id = (0..index.storage.num_partitions()) + .max_by_key(|partition| index.storage.partition_size(*partition)) + .expect("the index must hold at least one partition"); + let partition_rows = index.storage.partition_size(partition_id); + assert!( + partition_rows >= 20, + "partition {partition_id} holds {partition_rows} rows; the threshold is a \ + percentage, so a tiny partition makes both sides of it the same set" + ); + // Derive both sides from the actual partition size: cluster-to-partition + // assignment is not stable across IVF training changes, so fixed counts go + // stale. `few` sits strictly below the threshold, `many` at or above it. + let threshold_rows = (partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT).div_ceil(100); + let few: Vec = (0..threshold_rows.saturating_sub(1).max(1) as u32).collect(); + let many: Vec = (0..(threshold_rows + 1).min(partition_rows) as u32).collect(); + assert!( + many.len() * 100 >= partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT + && few.len() * 100 < partition_rows * COVERING_SCATTERED_READ_MAX_PERCENT, + "the fixture must put `few` below the threshold and `many` above it for a \ + {partition_rows}-row partition" + ); + let columns = vec!["payload".to_string()]; + + let read = async |positions: Option<&[u32]>| { + let io_stats = IoStats::new(); + let batch = index + .storage + .take_covering(partition_id, positions, &columns, Some(io_stats.clone())) + .await + .unwrap(); + (batch, io_stats.snapshot().bytes_read) + }; + + let (few_batch, few_bytes) = read(Some(&few)).await; + let (many_batch, many_bytes) = read(Some(&many)).await; + let (whole_batch, whole_bytes) = read(None).await; + + assert_eq!( + few_batch.num_rows(), + few.len(), + "below the threshold the gather must read only the survivors' rows" + ); + assert_eq!( + many_batch.num_rows(), + partition_rows, + "above the threshold the gather must fall back to the partition's whole range" + ); + assert_eq!( + whole_batch.num_rows(), + partition_rows, + "positions the caller could not derive read the whole range" + ); + + assert!(whole_bytes > 0, "the fixture must actually perform I/O"); + assert_eq!( + many_bytes, whole_bytes, + "the fallback is the same read as the whole-range one: {many_bytes} vs \ + {whole_bytes} bytes" + ); + assert!( + few_bytes * 4 < whole_bytes, + "the scattered read must be materially cheaper than the range it replaces: \ + {few_bytes} vs {whole_bytes} bytes for {} of {partition_rows} rows", + few.len() + ); + + // Neither branch is a stub: both return the right values for the rows asked for, + // matched by row id rather than by position. + let whole_row_ids = whole_batch + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let whole_payloads = whole_batch["payload"].as_string::(); + for batch in [&few_batch, &many_batch] { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let payloads = batch["payload"].as_string::(); + for (i, position) in few.iter().map(|p| *p as usize).enumerate() { + assert_eq!(row_ids.value(i), whole_row_ids.value(position)); + assert_eq!( + payloads.value(i), + whole_payloads.value(position), + "gathered payload does not belong to the row it was asked for" + ); + assert_eq!( + payloads.value(i), + format!("{:04}", row_ids.value(i)).repeat(256), + "gathered payload does not match the base table" + ); + } + } + } + + /// The multi-partition gather must hold `O(survivors)`, not one whole covering batch + /// per contributing partition -- even when every partition takes the whole-range + /// fallback, which is what a pending fragment-reuse index forces for all of them. + /// + /// This is a memory bound, so it is asserted against bytes measured in the same test + /// rather than a magic number: one partition's whole-range covering read is the + /// positive control, and the gather's result over *four* such partitions must stay + /// well under it. Row counts back it up -- a gather that accumulated whole partitions + /// returns every row of every one of them, not the four survivors. + /// + /// The line under test is the narrowing in `gather_survivor_covering` + /// (`take_record_batch` on each partition's batch before the next partition is read). + /// Pushing the untouched `gathered` batch instead fails both assertions here: + /// `4 * partition_rows` rows instead of 4, and a result larger than a whole partition. + #[tokio::test] + async fn test_covering_gather_holds_only_the_survivors_when_the_read_falls_back() { + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 256; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query = Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + 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(), + covering_projection: Some(Arc::from(vec!["payload".to_string()])), + }; + let covering = index + .query_covering(&query) + .unwrap() + .expect("the fixture declares covering columns"); + + // One survivor per partition, each with `position: None` -- the state a deferred + // fragment-reuse remap leaves every partition in, and the only one where the + // gather reads more than the survivors' own rows. + let mut locations: HashMap = HashMap::new(); + let mut expected_payloads: HashMap = HashMap::new(); + let mut whole_partition_bytes = 0; + for partition_id in 0..NUM_CLUSTERS { + let whole = index + .storage + .take_covering(partition_id, None, &covering.columns, None) + .await + .unwrap(); + assert_eq!( + whole.num_rows(), + index.storage.partition_size(partition_id), + "the control must read the partition's whole range" + ); + whole_partition_bytes = whole_partition_bytes.max(whole.get_array_memory_size()); + let row_id = whole + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::() + .value(0); + expected_payloads.insert(row_id, whole["payload"].as_string::().value(0).into()); + locations.insert( + row_id, + CoveringLocation { + partition_id, + position: None, + }, + ); + } + assert_eq!(locations.len(), NUM_CLUSTERS, "one survivor per partition"); + + let gathered = index + .gather_survivor_covering(&locations, &covering, None) + .await + .unwrap() + .expect("survivors were located, so the gather must return their covering"); + + assert_eq!( + gathered.num_rows(), + locations.len(), + "the gather must keep the survivors' rows only; holding every contributing \ + partition's whole covering would return {} rows", + (0..NUM_CLUSTERS) + .map(|p| index.storage.partition_size(p)) + .sum::() + ); + assert!( + whole_partition_bytes > 0, + "the fixture must actually hold covering bytes" + ); + assert!( + gathered.get_array_memory_size() * 4 < whole_partition_bytes, + "the gather must hold far less than a single partition's covering: {} bytes \ + for {} survivors vs {whole_partition_bytes} bytes for one partition", + gathered.get_array_memory_size(), + locations.len(), + ); + + // Bounded, and still correct: every survivor's own value, matched by row id. + let row_ids = gathered + .column_by_name(ROW_ID) + .expect("gathered batch carries row ids") + .as_primitive::(); + let payloads = gathered["payload"].as_string::(); + for (i, row_id) in row_ids.values().iter().enumerate() { + assert_eq!( + payloads.value(i), + expected_payloads[row_id], + "gathered payload does not belong to the row it was asked for" + ); + assert_eq!( + payloads.value(i), + format!("{row_id:04}").repeat(256), + "gathered payload does not match the base table" + ); + } + assert_eq!( + row_ids.values().iter().copied().collect::>(), + locations.keys().copied().collect::>(), + "the gather must return exactly the survivors it was asked for" + ); + } + + /// A covered query must return the covering values the index actually holds -- for a + /// **nullable** covering column too, where a null fill is indistinguishable from the + /// truth. + /// + /// This is the shape that hid a silent bug: when every covering fixture declares its + /// columns non-nullable, an implementation that emits nulls for the survivors fails + /// loudly on all of them and looks like a schema problem. A nullable column turns the + /// same defect into ten rows of silently wrong values, and only a comparison against + /// the base table catches it. `tag` is `1000 - offset` and `note` is `n{offset*7+3}`, + /// so neither equals its row id and a positional gather cannot coincide with a + /// by-name one. + /// + /// Both emit sites are covered: `query_parallelism = 1` merges partitions through the + /// global top-k heap (`global_heap_to_batch`), `> 1` searches each partition + /// separately (`append_covering`). Restoring only one of them fails one case. + /// + /// The absence of a `LanceRead` is what makes this an index measurement rather than a + /// scan measurement: a covering column is semantically transparent, so a plan that + /// re-fetched it from the base table would return byte-identical results. + #[rstest] + #[case::global_heap(1)] + #[case::per_partition(4)] + #[tokio::test] + async fn test_covered_query_returns_nullable_covering_values_from_the_index( + #[case] query_parallelism: i32, + ) { + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const K: usize = 10; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + + let query_key = vectors.value(0); + let mut scan = dataset.scan(); + scan.nearest("vector", query_key.as_primitive::(), K) + .unwrap(); + scan.minimum_nprobes(NUM_CLUSTERS); + scan.query_parallelism(query_parallelism); + scan.with_row_id(); + scan.project(&["tag", "note"]).unwrap(); + + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + !plan.contains("LanceRead"), + "a projection of covered columns only must not fetch from the base table, or \ + this test measures the scan rather than the index; plan:\n{plan}" + ); + assert!( + !plan.contains("payload"), + "the query reads neither `payload` nor anything derived from it, so the index \ + must not be asked to materialize it; plan:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), K, "k = {K} must return {K} rows"); + let row_ids = batch[ROW_ID].as_primitive::(); + let tags = batch["tag"].as_primitive::(); + let notes = batch["note"].as_string::(); + + let mut nulls = 0; + let mut values = 0; + for i in 0..batch.num_rows() { + let offset = row_ids.value(i) as usize; + assert_eq!( + tags.value(i), + 1000 - offset as i32, + "row {i} (row_id {offset}): covered tag belongs to another row" + ); + match offset % 3 { + 0 => { + assert!( + notes.is_null(i), + "row {i} (row_id {offset}): covered note must be NULL, got {:?}", + notes.value(i) + ); + nulls += 1; + } + _ => { + assert!( + !notes.is_null(i), + "row {i} (row_id {offset}): covered note must not be NULL -- a null \ + fill for every survivor is exactly the silent failure this fixture \ + exists to catch" + ); + assert_eq!( + notes.value(i), + format!("n{}", offset * 7 + 3), + "row {i} (row_id {offset}): covered note belongs to another row" + ); + values += 1; + } + } + } + assert!( + nulls > 0 && values > 0, + "the result must contain both NULL and non-NULL notes, or one of the two cases \ + above is vacuous (got {nulls} nulls, {values} values)" + ); + } + + /// `search_prepared_partition` is the synchronous phase of a prepared search: it runs + /// on the CPU pool, where the covering gather -- which is I/O -- cannot be awaited. + /// A query that needs covering columns must be told so, not handed rows whose covering + /// columns are quietly absent; the caller has async entry points for exactly this. + /// + /// The same call with `covering_projection = Some(&[])` must succeed and return the + /// bare `[_distance, _rowid]`. That is the state the projection narrowing exists for + /// and the one that silently degrades if it is folded into `None`: here it is the + /// difference between an error and a result. + /// + /// `supports_prepared_partition_search` must say so up front. A dispatcher that trusts + /// the flag pays for the partition load before it ever reaches the rejection, so a + /// covered index answering `true` there is a trap. The plain index built below is the + /// control: without it, `false` would be satisfied by an index that simply never + /// supported the entry point. + #[tokio::test] + async fn test_prepared_partition_search_rejects_a_query_that_needs_covering() { + use lance_index::prefilter::NoFilter; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, vectors) = covered_gather_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let query_for = |covering: Option>| Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: NUM_CLUSTERS, + 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(), + covering_projection: covering.map(Arc::from), + }; + + let prepared = index + .prepare_partition_search( + 0, + &query_for(Some(vec!["tag".to_string()])), + Arc::new(NoFilter), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let err = index + .search_prepared_partition(prepared, &NoOpMetricsCollector) + .expect_err("a covered query must not be silently served without its covering"); + assert!( + err.to_string().contains("search_prepared_partition"), + "the error must name the entry point that cannot serve it, got: {err}" + ); + + let prepared = index + .prepare_partition_search( + 0, + &query_for(Some(Vec::new())), + Arc::new(NoFilter), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + let batch = index + .search_prepared_partition(prepared, &NoOpMetricsCollector) + .expect("a query needing no covering column is served here as before"); + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + vec![DIST_COL, ROW_ID], + "`Some(&[])` means no covering work at all, not covering projected away" + ); + assert!(batch.num_rows() > 0, "the partition must be non-empty"); + + // The capability flag must not advertise what the rejection above denies. It + // carries no query, so it answers for the index: a covered index says no. + assert!( + !index.supports_prepared_partition_search(), + "a covered index must not advertise a prepared search it rejects" + ); + + // Control: the same index shape without covering columns still supports it. + let plain_dir = TempStrDir::default(); + let plain_schema = Arc::new(Schema::new(vec![Field::new( + "vector", + vectors.data_type().clone(), + false, + )])); + let plain_batch = + RecordBatch::try_new(plain_schema.clone(), vec![vectors.clone()]).unwrap(); + let mut plain_dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(plain_batch)], plain_schema), + plain_dir.as_str(), + None, + ) + .await + .unwrap(); + plain_dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &VectorIndexParams::ivf_flat(NUM_CLUSTERS, DistanceType::L2), + true, + ) + .await + .unwrap(); + let plain_ctx = load_vector_index_context(&plain_dataset, "vector", INDEX_NAME).await; + let plain_index = plain_ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + assert!( + plain_index.supports_prepared_partition_search(), + "an ordinary index still supports the prepared search; without this the \ + covered assertion above would pass on an index that never supported it" + ); + } + + /// A partition read with only its internal columns must still be a well-formed + /// storage: `try_from_batch` does not require covering columns, and `covering_batch()` + /// on the result reports none rather than erroring or fabricating an empty batch. + /// + /// The same partition read with `PartitionColumns::All` is the control. Without it, + /// "reports none" would pass just as well on an index that never had covering columns, + /// and the row-id comparison would have nothing to catch a narrowed read that silently + /// shifted rows. + #[tokio::test] + async fn test_codes_only_partition_storage_is_well_formed() { + use lance_index::vector::flat::storage::FLAT_COLUMN; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (dataset, _) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + // Whichever partition k-means actually filled, not partition 0 unconditionally: + // which centroid wins which rows depends on initialisation, and on aarch64 + // partition 0 came out empty and tripped the vacuity guard. Any non-empty + // partition exercises the same storage invariants, so this asserts on the + // storage layout rather than on the clustering. + let mut chosen = None; + for candidate in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(candidate, PartitionColumns::All, None) + .await + .unwrap(); + if !storage.is_empty() { + chosen = Some((candidate, storage)); + break; + } + } + let (partition_id, all) = + chosen.expect("some partition must be non-empty or every assertion below is vacuous"); + let internal = index + .load_partition_storage(partition_id, PartitionColumns::Internal, None) + .await + .unwrap(); + let covering_names = |storage: &lance_index::vector::flat::storage::FlatFloatStorage| { + let schema = storage.schema().clone(); + storage + .covering_field_indices() + .into_iter() + .map(|i| schema.field(i).name().to_string()) + .collect::>() + }; + assert_eq!( + covering_names(&all), + vec!["payload".to_string(), "price".to_string()], + "the control read must carry both covering columns, or the narrowed arm below \ + proves nothing" + ); + assert!(all.covering_batch().unwrap().is_some()); + + assert!( + covering_names(&internal).is_empty(), + "a codes-only read must carry no covering column" + ); + assert!( + internal.covering_batch().unwrap().is_none(), + "covering_batch() on a codes-only storage reports none rather than erroring \ + or emitting an empty batch" + ); + + // Well-formed: the storage's own columns are all there, and the rows are the same + // rows in the same order. Row ids are compared through `row_ids()`, which resolves + // `_rowid` by name -- the covered layout puts covering columns before it. + assert_eq!(internal.len(), all.len()); + assert_eq!( + internal.row_ids().copied().collect::>(), + all.row_ids().copied().collect::>(), + "narrowing the read must not disturb which rows the partition holds" + ); + assert!(internal.schema().column_with_name(ROW_ID).is_some()); + assert!(internal.schema().column_with_name(FLAT_COLUMN).is_some()); + assert_eq!( + internal.schema().fields().len(), + 2, + "flat storage's internal columns are exactly [{ROW_ID}, {FLAT_COLUMN}]; schema \ + was {:?}", + internal.schema() + ); + } + + /// The rebuild path reads a partition with `PartitionColumns::All` on purpose: it + /// re-writes those batches into the index being built, so a covering column left unread + /// is a covering column the merged index no longer has -- silently, since nothing + /// downstream of the rebuild asks for it again. + /// + /// Asserted on the merged STORAGE rather than through a query, so it holds independently + /// of what the search path currently does with covering. The line under test is + /// `PartitionColumns::All` in `IvfIndexBuilder::take_partition_batches`; switching it to + /// `Internal` drops the covering columns from the merged index and fails this test. + #[tokio::test] + async fn test_merge_optimize_preserves_covering_in_storage() { + use arrow_array::{Int32Array, StringArray, types::Int32Type}; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + const INDEXED: usize = NUM_CLUSTERS * ROWS_PER_CLUSTER; + const APPENDED: usize = 128; + const DIMS: usize = 16; + + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, _) = + covered_flat_fixture(test_uri, INDEX_NAME, NUM_CLUSTERS, ROWS_PER_CLUSTER).await; + + // Append unindexed rows, then merge them in. The merge is what re-reads the + // existing partitions and re-writes them into the new index file. + let mut flat = Vec::with_capacity(APPENDED * DIMS); + for row in 0..APPENDED { + for d in 0..DIMS { + flat.push(row as f32 * 0.01 + d as f32 * 0.0001); + } + } + let appended_vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(flat), DIMS as i32).unwrap(), + ); + let schema = dataset.schema().into(); + let appended = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from_iter_values( + (INDEXED..INDEXED + APPENDED).map(|i| -(i as i32) - 7), + )), + Arc::new(StringArray::from_iter_values( + (INDEXED..INDEXED + APPENDED).map(|i| format!("p{}", i * 3 + 11)), + )), + appended_vectors, + ], + ) + .unwrap(); + let appended_schema = appended.schema(); + dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(appended)], appended_schema), + test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .unwrap(); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let mut covered: HashMap = HashMap::new(); + for partition in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(partition, PartitionColumns::All, None) + .await + .unwrap(); + for batch in storage.to_batches().unwrap() { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("storage batch carries row ids") + .as_primitive::(); + let payloads = batch + .column_by_name("payload") + .expect("merged storage must keep covering column 'payload'") + .as_string::(); + let prices = batch + .column_by_name("price") + .expect("merged storage must keep covering column 'price'") + .as_primitive::(); + for i in 0..batch.num_rows() { + covered.insert( + row_ids.value(i), + (payloads.value(i).to_string(), prices.value(i)), + ); + } + } + } + assert_eq!( + covered.len(), + INDEXED + APPENDED, + "the merge must index every row, or the value check below misses the rows it dropped" + ); + + // Values, not just presence: a merge that re-attached covering by position rather + // than by row id keeps the column and corrupts it. + let row_ids: Vec = covered.keys().copied().collect(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["payload", "price"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + for (i, row_id) in row_ids.iter().enumerate() { + let (payload, price) = &covered[row_id]; + assert_eq!( + (payload.as_str(), *price), + (truth_payloads.value(i), truth_prices.value(i)), + "row_id {row_id}: merged covering value does not match the base table" + ); + } + } + + /// The remap path rewrites a partition into a new index file exactly as the merge + /// path does, so it has the same requirement: it must read every column, not just + /// the internal ones. It reaches the storage through the partition *entry* rather + /// than through `load_partition_storage`, which is why the merge test above does not + /// cover it -- and why a codes-only entry silently strips the covering columns from + /// every compacted covered index. + /// + /// Asserted on the remapped STORAGE, not through a query, so it holds independently + /// of what the search path currently does with covering (`test_ivf_pq_covered_survives_compaction` + /// asserts the same property through a query and is ignored until the survivors-only + /// gather lands). The line under test is `load_partition_entry_with_covering` in + /// `IvfIndexBuilder::remap`; the cached `load_partition` reads internal columns only + /// and fails this test. + #[tokio::test] + async fn test_compaction_remap_preserves_covering_in_storage() { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + use arrow_array::types::Int32Type; + + const INDEX_NAME: &str = "vector_idx"; + const NUM_CLUSTERS: usize = 4; + const ROWS_PER_CLUSTER: usize = 64; + + let test_dir = TempStrDir::default(); + let (mut dataset, _) = covered_flat_fixture( + test_dir.as_str(), + INDEX_NAME, + NUM_CLUSTERS, + ROWS_PER_CLUSTER, + ) + .await; + + // Delete enough rows to clear the materialize-deletions threshold, so compaction + // rewrites the fragment and remaps the index rather than no-opping. + dataset.delete("price > -60").await.unwrap(); + let metrics = compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + assert!( + metrics.files_removed > 0, + "compaction must actually rewrite fragments, or the remap under test never runs" + ); + + let ctx = load_vector_index_context(&dataset, "vector", INDEX_NAME).await; + let index = ctx + .index + .as_any() + .downcast_ref::() + .expect("expected IvfFlatIndex"); + + let mut covered: HashMap = HashMap::new(); + for partition in 0..index.ivf_model().num_partitions() { + let storage = index + .load_partition_storage(partition, PartitionColumns::All, None) + .await + .unwrap(); + for batch in storage.to_batches().unwrap() { + let row_ids = batch + .column_by_name(ROW_ID) + .expect("storage batch carries row ids") + .as_primitive::(); + let payloads = batch + .column_by_name("payload") + .expect("remapped storage must keep covering column 'payload'") + .as_string::(); + let prices = batch + .column_by_name("price") + .expect("remapped storage must keep covering column 'price'") + .as_primitive::(); + for i in 0..batch.num_rows() { + covered.insert( + row_ids.value(i), + (payloads.value(i).to_string(), prices.value(i)), + ); + } + } + } + assert!( + !covered.is_empty(), + "the remapped index must still hold rows, or the value check below is vacuous" + ); + + // Values, not just presence: a remap that re-attached covering by position rather + // than by row id keeps the column and corrupts it, and the row ids themselves are + // the post-compaction addresses. + let row_ids: Vec = covered.keys().copied().collect(); + let truth = dataset + .take_rows( + &row_ids, + crate::dataset::ProjectionRequest::from_columns( + ["payload", "price"], + dataset.schema(), + ), + ) + .await + .unwrap(); + let truth_payloads = truth.column_by_name("payload").unwrap().as_string::(); + let truth_prices = truth + .column_by_name("price") + .unwrap() + .as_primitive::(); + for (i, row_id) in row_ids.iter().enumerate() { + let (payload, price) = &covered[row_id]; + assert_eq!( + (payload.as_str(), *price), + (truth_payloads.value(i), truth_prices.value(i)), + "row_id {row_id}: remapped covering value does not match the base table" + ); + } + } + + /// 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. + /// + /// This is also the only test that reaches [`CoveringGather::WholeRange`]: the FRI + /// drops rows as the partition is loaded, so a storage position no longer addresses + /// the file and the survivor gather must re-read the whole range and match by row id. + /// Removing that alignment check fails here with `row id ... missing from covering + /// source` -- loudly, because the gather is matched by row id rather than by position. + #[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 @@ -6242,7 +8129,7 @@ mod tests { async fn load_partition_row_ids(index: &IvfPq, partition_idx: usize) -> Vec { index .storage - .load_partition(partition_idx, None) + .load_partition(partition_idx, PartitionColumns::Internal, None) .await .unwrap() .row_ids() @@ -6253,7 +8140,7 @@ mod tests { async fn load_flat_partition_row_ids(index: &IvfFlatIndex, partition_idx: usize) -> Vec { index .storage - .load_partition(partition_idx, None) + .load_partition(partition_idx, PartitionColumns::Internal, None) .await .unwrap() .row_ids() @@ -7775,6 +9662,214 @@ mod tests { ); } + #[tokio::test] + async fn test_optimize_refuses_to_restamp_rebound_covering() { + use crate::dataset::NewColumnTransform; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let dataset_uri = format!("{}/optimize_refuses_restamp", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let all_fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u32) + .collect(); + + 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()]); + + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(all_fragments) + .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 no fragment looks stale. + 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" + ); + + // The segment keeps its true build version, so the staleness pass runs and passes; + // only the id rebind is wrong. + for id in segment.fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + for id in segment.covering_fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![segment]) + .await + .unwrap(); + + // Before optimize the mismatch is visible and the read path does the right thing: + // the stamped source id disagrees with the declaration, so payload comes from the + // base table and is all null. + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "precondition: the rebound declaration must fall back to the base table" + ); + + // A rebuild copies that payload by name and would re-derive the stamp from the + // CURRENT schema, making the declaration and the stamp agree -- which would start + // serving the old field's values and elide the base-table read. Refuse instead. + let err = dataset + .optimize_indices(&OptimizeOptions::new()) + .await + .expect_err("optimize must not re-stamp a covering payload onto a different field id"); + let msg = err.to_string(); + assert!( + msg.contains("covering values stamped with source field ids") + && msg.contains(&original_field_id.to_string()) + && msg.contains(&new_field_id.to_string()), + "error must name both the stored and the current ids; got: {msg}" + ); + + // And the index is left in the state the read path already handles. + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "after the refusal the payload must still come from the base table" + ); + } + + /// The single-segment counterpart to the merge test above, and the case name and type + /// cannot see. A driver rebinds a segment built for a dropped `payload` onto a re-added + /// field with the same name and type but a fresh id. + /// + /// The declaration/payload contract permits the commit. At read time, however, the + /// storage's stamped source id must prevent those old values from being served under the + /// new field: planning falls back to a base-table take, whose re-added column is all null. + #[tokio::test] + async fn test_rebound_covering_field_id_falls_back_to_base_table() { + use crate::dataset::NewColumnTransform; + + let test_dir = TempStrDir::default(); + let base_uri = test_dir.as_str(); + let (schema, batches) = make_covered_test_batches(); + let query = batches[0]["vector"].as_fixed_size_list().value(0); + let dataset_uri = format!("{}/commit_rebind_covering_field_id", base_uri); + let mut dataset = write_dataset_from_batches(&dataset_uri, schema, batches).await; + + let all_fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u32) + .collect(); + + 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()]); + + let mut segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name("vec_idx".to_string()) + .fragments(all_fragments) + .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 no fragment looks stale. + 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" + ); + + // The segment keeps its true build version, so the staleness pass runs and passes; + // only the id rebind is wrong. + for id in segment.fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + for id in segment.covering_fields.iter_mut() { + if *id == original_field_id { + *id = new_field_id; + } + } + + dataset + .commit_existing_index_segments("vec_idx", "vector", vec![segment]) + .await + .unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vector", query.as_ref(), 10).unwrap(); + scan.nprobes(TWO_FRAG_NUM_PARTITIONS); + scan.project(&["payload"]).unwrap(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("LanceRead"), + "storage was built from field id {original_field_id}, not the declaration's \ + rebound id {new_field_id}, so payload must use a base-table take; plan:\n{plan}" + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 10); + let payload = batch["payload"].as_primitive::(); + assert_eq!( + payload.null_count(), + payload.len(), + "the re-added all-null payload must come from the base table, not stale index values" + ); + } + #[rstest] #[case::flat("IVF_HNSW_FLAT")] #[case::pq("IVF_HNSW_PQ")] @@ -8988,7 +11083,10 @@ mod tests { // Every partition, not just the first: a projection that applied // unevenly would leave the partition cache holding mixed schemas. for partition_id in 0..hnsw.ivf.num_partitions() { - let entry = hnsw.load_partition_entry(partition_id, None).await.unwrap(); + let entry = hnsw + .load_partition_entry(partition_id, PartitionColumns::Internal, None) + .await + .unwrap(); let loaded = entry.index.to_batch().unwrap(); let loaded_schema = loaded.schema(); let read = loaded_schema diff --git a/rust/lance/src/index/vector/pq.rs b/rust/lance/src/index/vector/pq.rs index 141c8b85f27..99280ff8e25 100644 --- a/rust/lance/src/index/vector/pq.rs +++ b/rust/lance/src/index/vector/pq.rs @@ -990,6 +990,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let is_empty_threads = Arc::new(Mutex::new(Vec::new())); let pre_filter = Arc::new(TestPreFilter::with_thread_capture( diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index b105cdfbefd..53a322572df 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -38,6 +38,30 @@ pub(super) fn gather_covering_columns_by_row_id( source: &RecordBatch, target_rowids: &arrow_array::UInt64Array, ) -> Result> { + let take_idx = row_id_take_indices(source, target_rowids)?; + 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) +} + +/// Where each of `target_rowids` sits in `source` (a batch carrying `_rowid`), as take +/// indices. Every target row id must be present; a miss is an error for the same reason +/// [`gather_covering_columns_by_row_id`] errors on one -- taking an unrelated row's +/// values would corrupt the covered result silently. +/// +/// Split out so a caller that must keep `_rowid` on the result (the survivor gather +/// narrows each partition's covering batch before reading the next one) can take the +/// whole batch instead of its non-`_rowid` columns. +pub(super) fn row_id_take_indices( + 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()))? @@ -47,7 +71,7 @@ pub(super) fn gather_covering_columns_by_row_id( for (i, rid) in src_rowids.values().iter().enumerate() { pos.insert(*rid, i as u32); } - let take_idx: arrow_array::UInt32Array = target_rowids + Ok(target_rowids .values() .iter() .map(|rid| { @@ -58,16 +82,7 @@ pub(super) fn gather_covering_columns_by_row_id( }) }) .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) + .into()) } /// Helper function to extract a column from a RecordBatch, supporting nested field paths. diff --git a/rust/lance/src/io/exec/ann_proto.rs b/rust/lance/src/io/exec/ann_proto.rs index 7ecb14211e6..7f749d6ad0d 100644 --- a/rust/lance/src/io/exec/ann_proto.rs +++ b/rust/lance/src/io/exec/ann_proto.rs @@ -25,7 +25,7 @@ use uuid::Uuid; use crate::Dataset; use crate::pb; -use super::knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec}; +use super::knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, resolve_physical_covering_query}; use super::table_identifier::{resolve_dataset, table_identifier_from_dataset}; use super::utils::PreFilterSource; @@ -118,9 +118,15 @@ pub fn query_to_proto(query: &Query) -> Result { dist_q_c: Some(query.dist_q_c), query_parallelism: Some(query.query_parallelism), approx_mode: approx_mode_to_proto(query.approx_mode) as i32, - // No planner narrows the covering projection yet, so this is always absent: - // "materialize every covering column declared". See `CoveringProjection`. - covering_projection: None, + // Three states, and the message wrapper is what keeps the middle one alive across + // the wire (see `CoveringProjection` in `ann.proto`). `map` -- never + // `unwrap_or_default()` -- because an empty list means "materialize nothing", + // while absence means "materialize everything the index declares". + covering_projection: query.covering_projection.as_ref().map(|columns| { + pb::CoveringProjection { + columns: columns.to_vec(), + } + }), }) } @@ -151,6 +157,13 @@ pub fn query_from_proto(proto: pb::VectorQueryProto) -> Result { query_parallelism: proto.query_parallelism.unwrap_or(DEFAULT_QUERY_PARALLELISM), dist_q_c: proto.dist_q_c.unwrap_or(0.0), approx_mode: approx_mode_from_proto(proto.approx_mode), + // A message field, so absence survives the round trip and stays distinct from an + // empty list. A producer predating field 15 sends nothing and the reconstructed + // query narrows nothing -- the conservative direction (see + // `Query::covering_projection`: under-narrowing costs time, never correctness). + covering_projection: proto + .covering_projection + .map(|projection| Arc::from(projection.columns)), }) } @@ -291,21 +304,31 @@ pub async fn ann_ivf_sub_index_exec_from_proto( } }; - ANNIvfSubIndexExec::try_new(input, dataset, indices, query, prefilter_source) + let query = resolve_physical_covering_query(dataset.clone(), &indices, &query).await?; + ANNIvfSubIndexExec::try_new_with_resolved_covering( + input, + dataset, + indices, + query, + prefilter_source, + ) } #[cfg(test)] mod tests { use super::*; - use arrow_array::types::{Float32Type, UInt32Type}; + use arrow_array::types::{Float32Type, Int32Type, UInt32Type}; use arrow_array::{ArrayRef, Float32Array, Float64Array}; use half::f16; use lance_datagen::{array, gen_batch}; use crate::index::DatasetIndexExt; use crate::index::vector::VectorIndexParams; + use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::test::TestMemoryExec; + use lance_core::ROW_ID; use lance_index::IndexType; + use lance_index::vector::DIST_COL; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; @@ -357,6 +380,7 @@ mod tests { query_parallelism: -1, dist_q_c: 0.42, approx_mode: ApproxMode::Accurate, + covering_projection: None, }; let proto = query_to_proto(&query).unwrap(); @@ -397,6 +421,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let proto = query_to_proto(&query).unwrap(); @@ -411,6 +436,148 @@ mod tests { assert_eq!(back.approx_mode, ApproxMode::Normal); } + /// `covering_projection` has three states and only two of them are expressible by a + /// bare proto3 `repeated string`, so the wire format wraps the list in a message. + /// Folding `Some(&[])` into `None` here would silently restore full materialization + /// on every distributed plan while every result stayed correct. + #[test] + fn test_query_roundtrip_covering_projection_three_states() { + let query = |projection: Option>| { + let key: ArrayRef = Arc::new(Float32Array::from(vec![1.0])); + Query { + column: "v".to_string(), + key, + k: 5, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: None, + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Normal, + covering_projection: projection.map(Arc::from), + } + }; + let roundtrip = |projection: Option>| { + let proto = query_to_proto(&query(projection)).unwrap(); + query_from_proto(proto).unwrap().covering_projection + }; + + assert!( + roundtrip(None).is_none(), + "'not computed' must stay 'not computed'" + ); + let narrowed = roundtrip(Some(vec![])).expect( + "'this query needs no covering column' must survive; decoding it as `None` \ + turns every remote plan back into a fully materializing one", + ); + assert!(narrowed.is_empty()); + assert_eq!( + roundtrip(Some(vec!["payload".to_string(), "price".to_string()])) + .expect("a computed projection must survive") + .to_vec(), + vec!["payload".to_string(), "price".to_string()], + ); + } + + /// The schema a remote executor declares must be the schema the planner declared. + /// The surrounding nodes -- `knn_combined`'s union and its projections -- were built + /// against the planner's, so a reconstruction that widens is not merely slower. + #[tokio::test] + async fn test_sub_index_proto_preserves_the_declared_covering_schema() { + let (dataset, _dir) = make_covered_indexed_dataset().await; + let indices = dataset.load_indices_by_name("vector_idx").await.unwrap(); + assert_eq!( + indices[0].covering_fields.len(), + 2, + "the fixture must actually be a covered index" + ); + + let input: Arc = + TestMemoryExec::try_new_exec( + &[], + super::super::knn::KNN_PARTITION_SCHEMA.clone(), + None, + ) + .unwrap(); + + for (projection, expected) in [ + ( + None, + vec![ + DIST_COL.to_string(), + ROW_ID.to_string(), + "price".to_string(), + "id".to_string(), + ], + ), + (Some(vec![]), vec![DIST_COL.to_string(), ROW_ID.to_string()]), + ( + Some(vec!["id".to_string()]), + vec![DIST_COL.to_string(), ROW_ID.to_string(), "id".to_string()], + ), + ] { + let key: ArrayRef = Arc::new(Float32Array::from(vec![0.1f32; 128])); + let query = Query { + column: "vector".to_string(), + key, + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 2, + 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: ApproxMode::Normal, + covering_projection: projection.clone().map(Arc::from), + }; + let query = resolve_physical_covering_query(dataset.clone(), &indices, &query) + .await + .unwrap(); + let exec = ANNIvfSubIndexExec::try_new_with_resolved_covering( + input.clone(), + dataset.clone(), + indices.clone(), + query, + PreFilterSource::None, + ) + .unwrap(); + let names = |schema: arrow_schema::SchemaRef| { + schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>() + }; + // Pin the local schema too: if this fixture stopped narrowing, the equality + // below would hold trivially for the wrong reason. + assert_eq!(names(exec.schema()), expected, "local plan, {projection:?}"); + + let proto = ann_ivf_sub_index_exec_to_proto(&exec).await.unwrap(); + let back = ann_ivf_sub_index_exec_from_proto( + proto, + Some(dataset.clone()), + input.clone(), + None, + ) + .await + .unwrap(); + assert_eq!( + names(back.schema()), + expected, + "reconstructed plan, {projection:?}" + ); + } + } + async fn make_vector_dataset() -> (Arc, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let batch = gen_batch() @@ -432,6 +599,41 @@ mod tests { (Arc::new(ds), dir) } + /// An IVF_PQ index covering `price` and `id`, declared in the reverse of schema + /// order so the assertions above can tell declaration order from schema order. + async fn make_covered_indexed_dataset() -> (Arc, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let batch = gen_batch() + .col("id", array::step::()) + .col("price", array::step::()) + .col( + "vector", + array::rand_vec::(lance_datagen::Dimension::from(128)), + ) + .into_batch_rows(lance_datagen::RowCount::from(256)) + .unwrap(); + let path = dir.path().join("test_ann_covered.lance"); + let mut ds = Dataset::write( + arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()), + path.to_str().unwrap(), + None, + ) + .await + .unwrap(); + + let mut index_params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::new(2), + PQBuildParams::default(), + ); + index_params.covering_columns(vec!["price".to_string(), "id".to_string()]); + ds.create_index(&["vector"], IndexType::Vector, None, &index_params, false) + .await + .unwrap(); + let ds = Dataset::open(path.to_str().unwrap()).await.unwrap(); + (Arc::new(ds), dir) + } + async fn make_indexed_dataset() -> (Arc, tempfile::TempDir) { let (_dataset, dir) = make_vector_dataset().await; let mut ds = Dataset::open(dir.path().join("test_ann.lance").to_str().unwrap()) @@ -475,6 +677,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let exec = ANNIvfPartitionExec::try_new( @@ -523,6 +726,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; // Use a TestMemoryExec as a mock input child (provides the KNN_PARTITION_SCHEMA) @@ -587,6 +791,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: ApproxMode::Normal, + covering_projection: None, }; let input: Arc = TestMemoryExec::try_new_exec( diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 75f97177de0..2a1ea633ef2 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -47,7 +47,7 @@ use lance_datafusion::utils::{ DELTAS_SEARCHED_METRIC, ExecutionPlanMetricsSetExt, FIND_PARTITIONS_ELAPSED_METRIC, PARTITIONS_RANKED_METRIC, PARTITIONS_SEARCHED_METRIC, }; -use lance_index::metrics::MetricsCollector; +use lance_index::metrics::{MetricsCollector, NoOpMetricsCollector}; use lance_index::prefilter::PreFilter; use lance_index::vector::DIST_Q_C_COLUMN; use lance_index::vector::{ @@ -64,6 +64,7 @@ use uuid::Uuid; use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder}; use crate::index::DatasetIndexInternalExt; +use crate::index::covering::{effective_covering, effective_covering_with_ids}; use crate::index::prefilter::{DatasetPreFilter, FilterLoader}; use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; use crate::{Error, Result}; @@ -1102,11 +1103,150 @@ pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { ])) }); +fn covering_fields_match(expected: &Field, physical: &Field) -> bool { + expected.name() == physical.name() + && expected.data_type() == physical.data_type() + && expected.is_nullable() == physical.is_nullable() + && expected.metadata() == physical.metadata() +} + +/// Intersect a requested manifest declaration with the capabilities proven by every +/// selected physical segment. +/// +/// The output remains in declaration order because that is the schema the ANN exec +/// exposes. If a segment stores the surviving columns in a different relative order then +/// none are served: storage emits in physical order, and pairing those values with a +/// declaration-ordered schema would be incorrect. +fn common_physical_covering( + mut requested: Vec<(i32, Field)>, + physical_by_segment: &[Vec<(i32, Field)>], +) -> Vec<(i32, Field)> { + if physical_by_segment.is_empty() { + return Vec::new(); + } + + for physical in physical_by_segment { + let before = requested.len(); + requested.retain(|(expected_id, expected_field)| { + let kept = physical.iter().any(|(physical_id, physical_field)| { + expected_id == physical_id && covering_fields_match(expected_field, physical_field) + }); + if !kept { + // Withdrawal is silent by design -- results stay correct, the plan just + // grows a base-table read -- which makes a covered index that quietly + // stopped eliding takes indistinguishable from one that never worked. + // `covering_fields_match` compares name, data type, nullability AND + // metadata, so the cause is usually a drift too small to guess at. + log::debug!( + "Covering column '{}' (field id {}) is declared but not servable by a \ + segment: no physical column matches it on name, type, nullability and \ + metadata. Its values will come from a base-table read.", + expected_field.name(), + expected_id + ); + } + kept + }); + if requested.is_empty() { + if before > 0 { + log::debug!( + "No declared covering column is servable by every selected segment; \ + this query reads all of them from the base table." + ); + } + return requested; + } + } + + let requested_ids = requested.iter().map(|(id, _)| *id).collect::>(); + for physical in physical_by_segment { + let physical_ids = physical + .iter() + .filter(|(physical_id, physical_field)| { + requested.iter().any(|(expected_id, expected_field)| { + expected_id == physical_id + && covering_fields_match(expected_field, physical_field) + }) + }) + .map(|(id, _)| *id) + .collect::>(); + if requested_ids != physical_ids { + // Segments disagree on the order or set they can serve. Storage emits in + // physical order while the schema is declaration-ordered, so pairing them + // would be wrong -- none are served rather than some. + log::debug!( + "Selected segments do not agree on the covering columns they can serve \ + ({requested_ids:?} vs {physical_ids:?}); this query reads all of them from \ + the base table." + ); + return Vec::new(); + } + } + requested +} + +/// Resolve the query's manifest-requested covering columns against every selected +/// segment's physical storage capability. +/// +/// The returned query always has an explicit covering projection. A missing physical +/// column is omitted from it, causing the normal scanner take to fetch that value from +/// the base table instead. +pub async fn resolve_physical_covering_query( + dataset: Arc, + indices: &[IndexMetadata], + query: &Query, +) -> Result { + let mut resolved_query = query.clone(); + let declared_ids = indices + .first() + .map(|index| index.covering_fields.as_slice()) + .unwrap_or_default(); + if let Some(mismatch) = indices + .iter() + .find(|index| index.covering_fields != declared_ids) + { + return Err(Error::index(format!( + "ANN index delta segments of index '{}' disagree on covering fields ({:?} vs {:?}); \ + the index metadata is inconsistent -- rebuild the index", + mismatch.name, declared_ids, mismatch.covering_fields + ))); + } + + let requested = effective_covering_with_ids( + declared_ids, + query.covering_projection.as_deref(), + dataset.schema(), + )?; + if requested.is_empty() { + resolved_query.covering_projection = Some(Arc::from(Vec::::new())); + return Ok(resolved_query); + } + + let physical_by_segment = future::try_join_all(indices.iter().map(|index| { + let dataset = dataset.clone(); + let column = query.column.clone(); + let uuid = index.uuid; + async move { + dataset + .open_vector_index(&column, &uuid, &NoOpMetricsCollector) + .await? + .physical_covering_fields() + } + })) + .await?; + let names = common_physical_covering(requested, &physical_by_segment) + .into_iter() + .map(|(_, field)| field.name().clone()) + .collect::>(); + resolved_query.covering_projection = Some(Arc::from(names)); + Ok(resolved_query) +} + /// Create a new ANN execution node. `overlay_block`, when `Some`, excludes rows whose index /// entries may be stale due to a newer data overlay (see [`ANNIvfSubIndexExec::with_overlay_block`]). /// `external_mask`, when `Some`, additionally restricts the scan to a caller-supplied /// allow/block set (see [`ANNIvfSubIndexExec::with_external_mask`]). -pub fn new_knn_exec( +pub async fn new_knn_exec( dataset: Arc, indices: &[IndexMetadata], query: &Query, @@ -1114,17 +1254,18 @@ pub fn new_knn_exec( overlay_block: Option, external_mask: Option>, ) -> Result> { + let query = resolve_physical_covering_query(dataset.clone(), indices, query).await?; let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), indices.iter().map(|idx| idx.uuid).collect_vec(), query.clone(), )?; - let mut sub_index = ANNIvfSubIndexExec::try_new( + let mut sub_index = ANNIvfSubIndexExec::try_new_with_resolved_covering( Arc::new(ivf_node), dataset, indices.to_vec(), - query.clone(), + query, prefilter_source, )?; if let Some(overlay_block) = overlay_block { @@ -1420,7 +1561,24 @@ pub struct ANNIvfSubIndexExec { } impl ANNIvfSubIndexExec { + /// Construct an ANN sub-index exec without a verified physical covering capability. + /// + /// This synchronous entry point cannot open segment storage, so it conservatively + /// disables covering and leaves declared values to the base-table take. Dataset + /// planning and protobuf reconstruction use the storage-aware async resolution path + /// before calling `Self::try_new_with_resolved_covering` (private, so not linked). pub fn try_new( + input: Arc, + dataset: Arc, + indices: Vec, + mut query: Query, + prefilter_source: PreFilterSource, + ) -> Result { + query.covering_projection = Some(Arc::from(Vec::::new())); + Self::try_new_with_resolved_covering(input, dataset, indices, query, prefilter_source) + } + + pub(crate) fn try_new_with_resolved_covering( input: Arc, dataset: Arc, indices: Vec, @@ -1433,11 +1591,9 @@ 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. + // Start from the logical declaration, narrowed to the explicit projection already + // proven against every selected segment's physical storage. The per-partition + // search emits only this subset, so TakeExec may skip the matching base-table fields. let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); let included_ids = indices .first() @@ -1458,33 +1614,21 @@ impl ANNIvfSubIndexExec { mismatch.name, included_ids, mismatch.covering_fields ))); } - 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(); + // Narrow only AFTER the agreement check above: narrowing first could reduce + // segments that genuinely disagree into a subset on which they happen to agree, + // masking metadata corruption as a working query. The manifest remains + // authoritative for logical identity and schema dependencies; the projection is + // authoritative for physical availability. + let covering = effective_covering( + &included_ids, + query.covering_projection.as_deref(), + dataset.schema(), + )?; + // Resolved-at-construction covered flag (see the `has_covered` field): from + // the NARROWED set, not the declaration -- a projection that needs none of + // the declared columns (`Some(&[])`) makes this plan non-covered. + let has_covered = !covering.is_empty(); + fields.extend(covering.into_iter().map(Arc::new)); let output_schema = Arc::new(arrow_schema::Schema::new(fields)); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), @@ -3001,13 +3145,71 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, } } - /// 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. + fn covering_field(id: i32, name: &str, data_type: DataType) -> (i32, Field) { + (id, Field::new(name, data_type, true)) + } + + #[test] + fn test_common_physical_covering_requires_every_segment_to_prove_the_field() { + let price = covering_field(17, "price", DataType::Int64); + let payload = covering_field(23, "payload", DataType::Utf8); + + let common = common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![price.clone(), payload.clone()], vec![price.clone()]], + ); + + assert_eq!(common, vec![price.clone()]); + + let common = common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![payload, price.clone()], vec![price.clone()]], + ); + assert_eq!( + common, + vec![price], + "order is checked after finding the plan-wide subset" + ); + } + + #[test] + fn test_common_physical_covering_rejects_wrong_identity_shape_and_order() { + let price = covering_field(17, "price", DataType::Int64); + let payload = covering_field(23, "payload", DataType::Utf8); + + assert!( + common_physical_covering( + vec![price.clone()], + &[vec![covering_field(99, "price", DataType::Int64)]], + ) + .is_empty(), + "a matching name and type cannot substitute for the source field id" + ); + assert!( + common_physical_covering( + vec![price.clone()], + &[vec![covering_field(17, "price", DataType::UInt64)]], + ) + .is_empty(), + "a matching source field id cannot substitute for the physical Arrow type" + ); + assert!( + common_physical_covering( + vec![price.clone(), payload.clone()], + &[vec![payload, price]], + ) + .is_empty(), + "a physical order that disagrees with the declared output schema is not servable" + ); + } + + /// The manifest remains authoritative for logical dependencies: if an index declares + /// a covering field id the current schema cannot resolve, construction must fail loudly + /// rather than hiding schema/metadata corruption behind physical fallback. #[tokio::test] async fn test_ann_sub_index_rejects_unresolvable_covering_field() { use lance_datafusion::exec::OneShotExec; @@ -3074,6 +3276,132 @@ mod tests { ); } + /// The exec's declared output schema is the seam through which a query's covering + /// projection reaches the rest of the plan: `knn_combined` unions the flat paths + /// against this schema, so whatever is declared here is what every unindexed + /// fragment pays to read. + /// + /// All three states of [`Query::covering_projection`] must be visible in the + /// declaration, and the `Some(&[])` case must NOT collapse into the `None` case. + /// The two are indistinguishable in query results -- a covering column is + /// semantically transparent, so the values simply arrive via a take instead -- which + /// is why this is asserted on the declared schema rather than on any result. + #[tokio::test] + async fn test_ann_sub_index_declares_only_the_projected_covering_columns() { + 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("vec", DataType::Int32, false), + ArrowField::new("price", DataType::Int32, true), + ArrowField::new("payload", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![Some(10), Some(20), Some(30)])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + "memory://d2-projected-covering", + None, + ) + .await + .unwrap(), + ); + let vec_id = dataset.schema().field("vec").unwrap().id; + let price_id = dataset.schema().field("price").unwrap().id; + let payload_id = dataset.schema().field("payload").unwrap().id; + + // Declared in the reverse of field-id order, so the assertions below distinguish + // declaration order from schema order. + let index = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![vec_id, payload_id, price_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![payload_id, price_id], + }; + index + .validate_covering_fields() + .expect("the fixture must be a legal covering declaration"); + + let declared_for = |projection: Option>| { + let mut query = base_query(); + query.covering_projection = projection.map(Arc::from); + let exec = ANNIvfSubIndexExec::try_new_with_resolved_covering( + input.clone(), + dataset.clone(), + vec![index.clone()], + query, + PreFilterSource::None, + ) + .unwrap(); + exec.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>() + }; + + assert_eq!( + declared_for(None), + vec![DIST_COL, ROW_ID, "payload", "price"], + "no projection means no narrowing: every declared covering column" + ); + assert_eq!( + declared_for(Some(vec!["price".to_string()])), + vec![DIST_COL, ROW_ID, "price"], + "only the covering column this query reads is declared" + ); + assert_eq!( + declared_for(Some(vec!["price".to_string(), "payload".to_string()])), + vec![DIST_COL, ROW_ID, "payload", "price"], + "covering columns are declared in the index's order, not the request's" + ); + assert_eq!( + declared_for(Some(vec![])), + vec![DIST_COL, ROW_ID], + "a query needing no covering column must declare a plain-index schema, which \ + is what lets the flat paths skip the covering read entirely" + ); + + let mut unverified_query = base_query(); + unverified_query.covering_projection = Some(Arc::from(["price".to_string()])); + let unverified = ANNIvfSubIndexExec::try_new( + input, + dataset, + vec![index], + unverified_query, + PreFilterSource::None, + ) + .unwrap(); + assert_eq!( + unverified + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![DIST_COL, ROW_ID], + "the synchronous constructor cannot verify storage and must disable covering" + ); + } + /// 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 @@ -4911,6 +5239,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; async fn multivector_scoring( @@ -4989,6 +5318,7 @@ mod tests { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, approx_mode: Default::default(), + covering_projection: None, }; let mut fields = KNN_INDEX_SCHEMA.fields().to_vec(); From 3734beda9a89d7326932a57bc047aeff504a80b7 Mon Sep 17 00:00:00 2001 From: Vivek Date: Mon, 17 Aug 2026 18:45:57 -0700 Subject: [PATCH 7/7] feat(bindings): build and inspect covered indexes from Java and Python Covering was reachable only from Rust: both bindings could read an index's `covering_fields` back, but neither could create a covered index, because both create paths passed an empty list. Java gains `VectorIndexParams.coveringColumns`, read through JNI into the Rust params. Python gains an `covering_columns` argument on `create_index`, `create_index_uncommitted` and the shared implementation they both funnel through -- all three, since threading only some of them leaves a path that silently builds an uncovered index. Passing it where it cannot be honoured now raises instead of being dropped. Python also surfaces covering when reading an index back: the ids an index carries, their names resolved against the current schema, and both in the description's repr, so a covered and an uncovered index no longer print identically. --- java/lance-jni/src/utils.rs | 13 +- java/src/main/java/org/lance/index/Index.java | 18 +- .../lance/index/vector/VectorIndexParams.java | 34 ++ .../java/org/lance/index/VectorIndexTest.java | 358 +++++++++++++++++- python/python/lance/dataset.py | 25 ++ .../python/lance/lance/indices/__init__.pyi | 2 + python/python/tests/test_commit_index.py | 71 +++- python/python/tests/test_vector_index.py | 256 +++++++++++++ python/src/dataset.rs | 10 + python/src/indices.rs | 204 +++++++++- 10 files changed, 981 insertions(+), 10 deletions(-) diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 29feff8f419..de7ec95ea84 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -24,7 +24,7 @@ use lance_linalg::distance::DistanceType; use crate::error::{Error, Result}; use crate::ffi::JNIEnvExt; -use crate::traits::FromJObjectWithEnv; +use crate::traits::{FromJObjectWithEnv, import_vec_from_method}; use lance_index::vector::{ApproxMode, Query}; use std::collections::HashMap; use std::str::FromStr; @@ -519,13 +519,22 @@ pub fn get_vector_index_params( stages.push(StageParams::RQ(rq_params)); } + // Covering ("included") columns: names of extra dataset columns stored inline in + // the index. Empty when absent. The core validates them when the index is built. + let covering_columns: Vec = import_vec_from_method( + env, + &vector_index_params_obj, + "getCoveringColumns", + |env, elem| Ok(env.get_string(&JString::from(elem))?.into()), + )?; + Ok(VectorIndexParams { metric_type: distance_type, stages, version: IndexFileVersion::V3, skip_transpose: false, runtime_hints: Default::default(), - covering_columns: Default::default(), + covering_columns, }) }, )?; diff --git a/java/src/main/java/org/lance/index/Index.java b/java/src/main/java/org/lance/index/Index.java index 620474ff787..bf4b1cf4418 100644 --- a/java/src/main/java/org/lance/index/Index.java +++ b/java/src/main/java/org/lance/index/Index.java @@ -16,6 +16,7 @@ import com.google.common.base.MoreObjects; import java.time.Instant; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -55,8 +56,19 @@ private Index( Long sizeBytes, IndexType indexType) { this.uuid = uuid; - this.fields = fields; - this.coveringFields = coveringFields; + // Both lists are defensively copied and unmodifiable so a caller cannot mutate them + // after construction: equals/hashCode would silently drift from the index files that + // actually carry the payload, and `coveringFields` is documented as the trailing slice + // of `fields`, so mutating either alone breaks that relationship. Null-coalesced so an + // explicit null from the builder reads as "none" rather than blowing up at first use. + this.fields = + fields == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(fields)); + this.coveringFields = + coveringFields == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(coveringFields)); this.name = name; this.datasetVersion = datasetVersion; this.fragments = fragments; @@ -89,7 +101,7 @@ public List fields() { *

These ids also appear in {@link #fields()} — that is deliberate, so that every consumer * reading {@code fields()} as the index's dependency set also covers them with no change. * - * @return the covering field IDs + * @return the covering field IDs, as an unmodifiable list */ public List coveringFields() { return coveringFields; diff --git a/java/src/main/java/org/lance/index/vector/VectorIndexParams.java b/java/src/main/java/org/lance/index/vector/VectorIndexParams.java index e8928943b48..f0faa46b1cb 100644 --- a/java/src/main/java/org/lance/index/vector/VectorIndexParams.java +++ b/java/src/main/java/org/lance/index/vector/VectorIndexParams.java @@ -17,6 +17,9 @@ import com.google.common.base.MoreObjects; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Optional; /** Parameters for creating a vector index. */ @@ -27,6 +30,7 @@ public class VectorIndexParams { private final Optional hnswParams; private final Optional sqParams; private final Optional rqParams; + private final List coveringColumns; private VectorIndexParams(Builder builder) { this.distanceType = builder.distanceType; @@ -35,6 +39,10 @@ private VectorIndexParams(Builder builder) { this.hnswParams = builder.hnswParams; this.sqParams = builder.sqParams; this.rqParams = builder.rqParams; + this.coveringColumns = + builder.coveringColumns == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(builder.coveringColumns)); validate(); } @@ -179,6 +187,7 @@ public static class Builder { private Optional hnswParams = Optional.empty(); private Optional sqParams = Optional.empty(); private Optional rqParams = Optional.empty(); + private List coveringColumns = Collections.emptyList(); /** * Create a new builder to create a vector index. @@ -235,6 +244,20 @@ public Builder setRqParams(RQBuildParams rqParams) { return this; } + /** + * Set the columns to cover ("include") in the index. Their values are stored inline in the + * index so a query projecting only covered columns is answered from the index without a take + * from the base table. Each must name a top-level, non-key column of the dataset; the columns + * are validated by the core when the index is built. Empty by default (no covering). + * + * @param coveringColumns the covering column names + * @return Builder + */ + public Builder setCoveringColumns(List coveringColumns) { + this.coveringColumns = coveringColumns; + return this; + } + public VectorIndexParams build() { return new VectorIndexParams(this); } @@ -268,6 +291,16 @@ public Optional getRqParams() { return rqParams; } + /** + * Get the covering ("included") columns whose values are stored inline in the index. Empty when + * the index has no covering columns. + * + * @return the covering column names, as an unmodifiable list + */ + public List getCoveringColumns() { + return coveringColumns; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -277,6 +310,7 @@ public String toString() { .add("hnswParams", hnswParams.orElse(null)) .add("sqParams", sqParams.orElse(null)) .add("rqParams", rqParams.orElse(null)) + .add("coveringColumns", coveringColumns) .toString(); } } diff --git a/java/src/test/java/org/lance/index/VectorIndexTest.java b/java/src/test/java/org/lance/index/VectorIndexTest.java index 7a9e066a1c5..e6f6bf7c03e 100755 --- a/java/src/test/java/org/lance/index/VectorIndexTest.java +++ b/java/src/test/java/org/lance/index/VectorIndexTest.java @@ -22,20 +22,36 @@ import org.lance.index.vector.SQBuildParams; import org.lance.index.vector.VectorIndexParams; import org.lance.index.vector.VectorTrainer; - +import org.lance.ipc.Query; +import org.lance.ipc.ScanOptions; + +import org.apache.arrow.dataset.scanner.Scanner; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; +import java.nio.charset.StandardCharsets; import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; +import java.util.Set; +import java.util.UUID; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class VectorIndexTest { @@ -347,6 +363,346 @@ public void testCreateIvfRqIndex(@TempDir Path tempDir) throws Exception { } } + /** + * {@code coveringColumns} must be a defensive, unmodifiable copy: mutating the list the caller + * passed to the builder (or the list the getter returns) must not change the params, so the JNI + * read at build time sees exactly what was requested. + */ + @Test + public void testCoveringColumnsIsDefensivelyCopied() { + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(2).build(); + List cols = new ArrayList<>(); + cols.add("i"); + VectorIndexParams params = new VectorIndexParams.Builder(ivf).setCoveringColumns(cols).build(); + + // Mutating the caller's list after build must not affect the params. + cols.add("s"); + cols.clear(); + assertEquals( + Collections.singletonList("i"), + params.getCoveringColumns(), + "include columns must be a defensive copy, unaffected by caller mutation"); + + // The returned list must be unmodifiable. + assertThrows(UnsupportedOperationException.class, () -> params.getCoveringColumns().add("x")); + } + + /** + * {@code Index.coveringFields} must be a defensive, unmodifiable copy, matching {@code + * VectorIndexParams.coveringColumns}: mutating the caller's list (or the list the getter returns) + * must not change the value object -- equals/hashCode would silently drift, and a later JNI + * commit would read mutated covering metadata while the index files still carry the payload. + */ + @Test + public void testIndexCoveringFieldsIsDefensivelyCopied() { + List ids = new ArrayList<>(); + ids.add(3); + Index index = + Index.builder() + .uuid(UUID.randomUUID()) + .fields(Collections.singletonList(0)) + .name("covered_idx") + .datasetVersion(1L) + .indexVersion(0) + .coveringFields(ids) + .build(); + + // Mutating the caller's list after build must not affect the index. + ids.add(7); + ids.clear(); + assertEquals( + Collections.singletonList(3), + index.coveringFields(), + "covering fields must be a defensive copy, unaffected by caller mutation"); + + // The returned list must be unmodifiable. + assertThrows(UnsupportedOperationException.class, () -> index.coveringFields().add(9)); + } + + /** + * A covered ("included") column set on {@link VectorIndexParams} must be threaded through the JNI + * create path into the built index's metadata, so the committed index reports the covered field + * id. Without the wiring the index would build but silently carry no covering columns. + * + *

This guards the {@code covering_columns} read in {@code java/lance-jni/src/utils.rs}: + * restoring it to an empty default leaves the index buildable but drops the covering declaration, + * and the field-id assertions below fail. + */ + @Test + public void testCreateIvfFlatIndexWithCoveringColumns(@TempDir Path tempDir) throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("ivf_flat_covering"))) { + try (Dataset dataset = testVectorDataset.create()) { + // Several partitions, so the search below really probes an IVF partition instead of + // degenerating into a scan of the one partition that holds everything. + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(4).build(); + + // Cover the non-vector "i" column so a projection of it is answered from the index. + VectorIndexParams vectorIndexParams = + new VectorIndexParams.Builder(ivf) + .setDistanceType(DistanceType.L2) + .setCoveringColumns(Collections.singletonList("i")) + .build(); + IndexParams indexParams = + IndexParams.builder().setVectorIndexParams(vectorIndexParams).build(); + + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(TestVectorDataset.vectorColumnName), + IndexType.IVF_FLAT, + indexParams) + .withIndexName(TestVectorDataset.indexName) + .build()); + + Index covered = + dataset.getIndexes().stream() + .filter(idx -> TestVectorDataset.indexName.equals(idx.name())) + .findFirst() + .orElse(null); + assertNotNull(covered, "Expected covered IVF_FLAT index to be present"); + + int coveredFieldId = fieldId(dataset, "i"); + int vectorFieldId = fieldId(dataset, TestVectorDataset.vectorColumnName); + assertEquals( + Collections.singletonList(coveredFieldId), + covered.coveringFields(), + "committed index must report the requested covering column's field id"); + // covering fields are always the trailing entries of fields, so the keyed vector field + // comes first and the covering field is appended after it. + assertEquals( + Arrays.asList(vectorFieldId, coveredFieldId), + covered.fields(), + "covering field must be appended after the keyed vector field"); + + float[] key = new float[32]; + for (int j = 0; j < key.length; j++) { + key[j] = (float) (32 + j); + } + int k = 5; + + // The assertion that actually exercises the covering payload: project the covered + // column together with an *uncovered* one and cross-check them row by row. Every row + // of the fixture satisfies s == "s-" + i, and only "i" is carried by the index, so "s" + // necessarily comes from a base-table take. A covering payload read positionally + // rather than by name, or attached to the wrong row, makes the two disagree. + Map coveredRows = searchCoveredWithBaseColumn(dataset, key, k); + assertEquals(k, coveredRows.size(), "covered search should return k distinct rows"); + coveredRows.forEach( + (i, s) -> + assertEquals( + "s-" + i, + s, + "covered 'i' must belong to the same row as the base-table 's'; a" + + " misaligned covering payload disagrees here")); + + // Recall against an exact scan. Worth keeping as a floor on search quality, but note + // what it cannot show: it does NOT establish that the *index* served the query. + // Covering is semantically transparent -- if the index were ignored and Lance fell + // back to brute-force KNN, the fallback is exact and would score 1.0 here. Java has + // no explain-plan binding, so there is no node chain to assert against as + // test_create_index_covering_columns_serve_the_query_from_the_index does on the + // Python side. That distinction is pinned instead by + // testCoveredIndexServesTheSearchNotAFullScan below, which uses fast search over a + // partially indexed dataset so the two outcomes genuinely differ. + Set exact = searchCoveredColumn(dataset, key, k, false); + Set ann = searchCoveredColumn(dataset, key, k, true); + assertEquals(k, exact.size(), "exact KNN ground truth should return k distinct rows"); + Set hits = new HashSet<>(ann); + hits.retainAll(exact); + double recall = (double) hits.size() / k; + assertTrue( + recall >= 0.5, + "recall of the covered index against exact KNN must be at least 0.5 but was " + + recall + + " (ann=" + + ann + + ", exact=" + + exact + + ")"); + } + } + } + + /** + * A covered search must actually be served by the index, not by a full scan that silently + * produces the same answer. + * + *

Recall against an exact scan cannot show this: covering is semantically transparent, so if + * the index became unusable and Lance fell back to brute-force KNN, the fallback is exact and + * scores 1.0. Java has no explain-plan binding, so the Python trick of pinning the plan's node + * chain is unavailable. + * + *

What is available is fast search, which restricts a query to indexed fragments. Index only 2 + * of the dataset's 5 fragments, and the two outcomes stop agreeing: {@code TestVectorDataset} + * writes the *same* 80 vectors into every fragment, so the query key matches one row at distance + * 0 in each of the 5 fragments ("i" = 1, 81, 161, 241, 321) with the next-nearest row ~181 away. + * Served by the index under fast search, only the two indexed fragments can contribute, so every + * returned "i" is below 160. Served by a full scan, the distance-0 rows from the three unindexed + * fragments win and "i" values above 160 appear. + */ + @Test + public void testCoveredIndexServesTheSearchNotAFullScan(@TempDir Path tempDir) throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("ivf_flat_covering_fast_search"))) { + try (Dataset dataset = testVectorDataset.create()) { + List fragments = dataset.getFragments(); + assertTrue(fragments.size() >= 3, "fixture must have unindexed fragments left over"); + + IvfBuildParams ivf = new IvfBuildParams.Builder().setNumPartitions(4).build(); + VectorIndexParams vectorIndexParams = + new VectorIndexParams.Builder(ivf) + .setDistanceType(DistanceType.L2) + .setCoveringColumns(Collections.singletonList("i")) + .build(); + IndexParams indexParams = + IndexParams.builder().setVectorIndexParams(vectorIndexParams).build(); + + List segments = new ArrayList<>(); + for (int f = 0; f < 2; f++) { + segments.add( + dataset.createIndex( + IndexOptions.builder( + Collections.singletonList(TestVectorDataset.vectorColumnName), + IndexType.IVF_FLAT, + indexParams) + .withIndexName(TestVectorDataset.indexName) + .withFragmentIds(Collections.singletonList(fragments.get(f).getId())) + .build())); + } + dataset.commitExistingIndexSegments( + TestVectorDataset.indexName, TestVectorDataset.vectorColumnName, segments); + + int coveredFieldId = fieldId(dataset, "i"); + for (Index segment : segments) { + assertEquals( + Collections.singletonList(coveredFieldId), + segment.coveringFields(), + "every distributed segment must carry the covering declaration"); + } + + float[] key = new float[32]; + for (int j = 0; j < key.length; j++) { + key[j] = (float) (32 + j); + } + + ScanOptions options = + new ScanOptions.Builder() + .columns(Collections.singletonList("i")) + .fastSearch(true) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(5) + .setUseIndex(true) + .build()) + .build(); + + Set values = new HashSet<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + for (int row = 0; row < root.getRowCount(); row++) { + values.add(iVector.get(row)); + } + } + } + + assertFalse(values.isEmpty(), "fast search over the covered index returned no rows"); + // Positive: the exact-match rows of both indexed fragments must be present. + assertTrue( + values.contains(1) && values.contains(81), + "fast search must return the distance-0 row of each indexed fragment, got " + values); + // Discriminating: any row from an unindexed fragment means a full scan served this. + for (int value : values) { + assertTrue( + value < 160, + "fast search must be confined to the 2 indexed fragments, but returned i=" + + value + + " (a full scan, not the covered index, served this query); got " + + values); + } + } + } + } + + /** Resolve a top-level column's Lance field id, which is what index metadata records. */ + private static int fieldId(Dataset dataset, String name) { + return dataset.getLanceSchema().fields().stream() + .filter(field -> name.equals(field.getName())) + .findFirst() + .orElseThrow(() -> new AssertionError("No field named " + name + " in the dataset schema")) + .getId(); + } + + /** + * Run a nearest-neighbor search projecting the covered "i" column together with the uncovered "s" + * column, returning i -> s for the matched rows. Only "i" is carried by the index, so "s" + * comes from a base-table take and the pair cross-checks the covering payload's row alignment. + */ + private static Map searchCoveredWithBaseColumn( + Dataset dataset, float[] key, int k) throws Exception { + ScanOptions options = + new ScanOptions.Builder() + .columns(Arrays.asList("i", "s")) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(k) + .setUseIndex(true) + .build()) + .build(); + + Map rows = new HashMap<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + VarCharVector sVector = (VarCharVector) root.getVector("s"); + for (int row = 0; row < root.getRowCount(); row++) { + rows.put(iVector.get(row), new String(sVector.get(row), StandardCharsets.UTF_8)); + } + } + } + return rows; + } + + /** + * Run a nearest-neighbor search projecting only the covered "i" column and return the matched + * values, so a covered result can be compared against exact KNN ground truth. + */ + private static Set searchCoveredColumn( + Dataset dataset, float[] key, int k, boolean useIndex) throws Exception { + ScanOptions options = + new ScanOptions.Builder() + .columns(Collections.singletonList("i")) + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(k) + .setUseIndex(useIndex) + .build()) + .build(); + + Set values = new HashSet<>(); + try (Scanner scanner = dataset.newScan(options); + ArrowReader reader = scanner.scanBatches()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + IntVector iVector = (IntVector) root.getVector("i"); + for (int row = 0; row < root.getRowCount(); row++) { + values.add(iVector.get(row)); + } + } + } + return values; + } + /** * Regression test for the metric_type passthrough in the JNI VectorTrainer. * diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 616c5f31004..fc0ea39b8a0 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3848,6 +3848,7 @@ def _create_index_impl( streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, rabitq_model: Optional[str] = None, + covering_columns: Optional[List[str]] = None, require_commit: bool = True, **kwargs, ) -> Index: @@ -4176,6 +4177,9 @@ def _create_index_impl( if rabitq_model is not None: kwargs["rabitq_model"] = rabitq_model + if covering_columns is not None: + kwargs["covering_columns"] = covering_columns + # Add fragment_ids and index_uuid to kwargs if provided for # distributed indexing if fragment_ids is not None: @@ -4236,6 +4240,7 @@ def create_index( streaming_coreset_rate: Optional[int] = None, streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, + covering_columns: Optional[List[str]] = None, progress_callback: Optional[Callable[[IndexProgress], None]] = None, **kwargs, ) -> LanceDataset: @@ -4305,6 +4310,13 @@ def create_index( If True, the index will be trained on the data (e.g., compute IVF centroids, PQ codebooks). If False, an empty index structure will be created without training, which can be populated later. + covering_columns : List[str], optional + The columns to cover ("include") in the index. Their values are stored + inline in the index so a query projecting only covered columns is + answered from the index without a take from the base table. Each must + name a top-level, non-key column of the dataset; the columns are + validated by the core when the index is built. Vector indexes only. + None by default (no covering). fragment_ids : List[int], optional If provided, the index will be created only on the specified fragments. This enables distributed/fragment-level indexing. When provided, the @@ -4471,6 +4483,7 @@ def create_index( streaming_coreset_rate=streaming_coreset_rate, streaming_refine_passes=streaming_refine_passes, skip_transpose=skip_transpose, + covering_columns=covering_columns, require_commit=True, **kwargs, ) @@ -4509,6 +4522,7 @@ def create_index_uncommitted( streaming_refine_passes: Optional[int] = None, skip_transpose: bool = False, rabitq_model: Optional[str] = None, + covering_columns: Optional[List[str]] = None, **kwargs, ) -> Index: """ @@ -4571,6 +4585,16 @@ def create_index_uncommitted( "create_index_uncommitted requires fragment_ids " "for distributed index build" ) + # Covering columns are a vector-index feature. This branch builds a + # scalar segment and never reaches the vector params, so honouring + # covering_columns here is impossible -- reject it rather than return a + # segment with covering_fields == [] that the caller believes is covered. + if covering_columns: + raise ValueError( + "covering_columns is only supported for vector indexes, but " + f"index_type={index_type!r} builds a scalar index segment; " + f"got covering_columns={covering_columns!r}" + ) kwargs = dict(kwargs) column, rust_index_type, _ = self._prepare_scalar_index_request( @@ -4617,6 +4641,7 @@ def create_index_uncommitted( streaming_refine_passes=streaming_refine_passes, skip_transpose=skip_transpose, rabitq_model=rabitq_model, + covering_columns=covering_columns, require_commit=False, **kwargs, ) diff --git a/python/python/lance/lance/indices/__init__.pyi b/python/python/lance/lance/indices/__init__.pyi index f4b0bf69592..915844774f7 100644 --- a/python/python/lance/lance/indices/__init__.pyi +++ b/python/python/lance/lance/indices/__init__.pyi @@ -92,6 +92,8 @@ class IndexDescription: num_rows_indexed: int fields: list[int] field_names: list[str] + covering_fields: list[int] + covering_field_names: list[str] segments: list[IndexSegmentDescription] details: dict total_size_bytes: Optional[int] diff --git a/python/python/tests/test_commit_index.py b/python/python/tests/test_commit_index.py index f4939434957..db8bbc16633 100644 --- a/python/python/tests/test_commit_index.py +++ b/python/python/tests/test_commit_index.py @@ -279,8 +279,23 @@ def test_index_covering_fields_roundtrip(dataset_with_index, test_table, tmp_pat # Exercises indices.rs: the IndexMetadata -> PyIndexSegmentDescription # conversion used by describe_indices(). - segment = dataset_without_index.describe_indices()[0].segments[0] - assert segment.covering_fields == [price_id] + desc = dataset_without_index.describe_indices()[0] + assert desc.segments[0].covering_fields == [price_id] + + # The index-level view lifts the same declaration out of the segments and + # resolves it against the current schema. + assert desc.covering_fields == [price_id] + assert desc.covering_field_names == ["price"] + # The covering suffix is reported apart from the key: `fields` stays the + # keyed prefix so it never claims the index can answer queries on "price". + assert desc.fields == [meta_id] + assert desc.field_names == ["meta"] + + # The repr has to carry the covering too, or `print(ds.describe_indices()[0])` + # renders a covered and an uncovered index identically. + rendered = repr(desc) + assert f"covering_fields=[{price_id}]" in rendered + assert 'covering_field_names=["price"]' in rendered def test_commit_index_rejects_invalid_covering_fields(dataset_with_index, tmp_path): @@ -316,3 +331,55 @@ def test_commit_index_rejects_invalid_covering_fields(dataset_with_index, tmp_pa create_index_op, read_version=dataset_with_index.version, ) + + +def test_commit_index_tolerates_missing_covering_fields( + dataset_with_index, test_table, tmp_path +): + """An Index pickled by a pre-covering lance has no `covering_fields` attribute. + + The dataclass default is `field(default_factory=list)`, so pickle's + `__dict__`-only restore leaves nothing to fall back to -- the attribute is + simply absent. The PyO3 conversion must treat that as "uncovered" (empty), + like it already does for `index_details`, not abort the commit with an + AttributeError mid rolling upgrade. + """ + from lance.dataset import Index + + index_id = dataset_with_index.describe_indices()[0].segments[0].uuid + + dataset_without_index = lance.write_dataset( + test_table, tmp_path / "dataset_without_index" + ) + src_index_dir = Path(dataset_with_index.uri) / "_indices" / index_id + dest_index_dir = Path(dataset_without_index.uri) / "_indices" / index_id + shutil.copytree(src_index_dir, dest_index_dir) + + field_id = _get_field_id_by_name(dataset_without_index.lance_schema, "meta") + index = Index( + uuid=index_id, + name="meta_idx", + fields=[field_id], + dataset_version=dataset_without_index.version, + fragment_ids=set( + [f.fragment_id for f in dataset_without_index.get_fragments()] + ), + index_version=0, + ) + # Simulate the old-writer object: the attribute does not exist at all, + # which is different from `covering_fields=[]`. + delattr(index, "covering_fields") + assert not hasattr(index, "covering_fields") + + create_index_op = lance.LanceOperation.CreateIndex( + new_indices=[index], + removed_indices=[], + ) + committed = lance.LanceDataset.commit( + dataset_without_index.uri, + create_index_op, + read_version=dataset_without_index.version, + ) + + described = committed.describe_indices()[0] + assert described.segments[0].covering_fields == [] diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 253592dd35b..35a4a52a999 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -328,6 +328,262 @@ def test_ann(indexed_dataset): run(indexed_dataset) +def _field_id(dataset, name): + for field in dataset.lance_schema.fields(): + if field.name() == name: + return field.id() + raise KeyError(name) + + +def test_create_index_with_covering_columns(tmp_path): + """A covered ("included") column passed to create_index is threaded through the + PyO3 boundary into the built index's metadata, so the committed index reports the + covered field id. Without the wiring the index builds but carries no covering + columns.""" + tbl = create_table() + dataset = lance.write_dataset(tbl, tmp_path) + price_id = _field_id(dataset, "price") + + query = np.random.randn(128).astype(np.float32) + exact = dataset.to_table( + nearest={"column": "vector", "q": query, "k": 10}, columns=["id"] + )["id"].to_pylist() + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["price"], + ) + + # The CreateIndex transaction carries the built index's metadata, including the + # covered field ids. + created_index = dataset.get_transactions(1)[0].operation.new_indices[0] + assert created_index.covering_fields == [price_id] + + # The covering declaration also survives into the committed manifest, so it is + # visible without replaying the transaction log. + desc = dataset.describe_indices()[0] + assert [segment.covering_fields for segment in desc.segments] == [[price_id]] + + # The index-level view resolves the same declaration to schema names, so callers + # can see what an index covers without knowing field ids. + assert desc.covering_fields == [price_id] + assert desc.covering_field_names == ["price"] + # Covering columns are reported separately from the key the index answers on. + assert desc.field_names == ["vector"] + + hits = dataset.to_table( + nearest={ + "column": "vector", + "q": query, + "k": 10, + "nprobes": 4, + "refine_factor": 10, + }, + columns=["id", "price"], + ) + recall = len(set(exact) & set(hits["id"].to_pylist())) / len(exact) + assert recall >= 0.5, f"recall={recall}, exact={exact}, got={hits['id']}" + + # Declaring a covering column must not disturb the values a query projects. + prices = dict(zip(tbl["id"].to_pylist(), tbl["price"].to_pylist())) + assert hits["price"].to_pylist() == [prices[i] for i in hits["id"].to_pylist()] + + +def _covered_query_table(nvec=1000, ndim=128, k=10, id_offset=1_000_000, seed=0): + """A vector table with an unambiguous k-neighborhood and ids disjoint from row ids. + + Two fixture properties are load-bearing, and neither is cosmetic: + + `create_table` numbers ids from 0, so on a single-fragment write `id == _rowid` for + every row -- and under that fixture any defect that serves the row id where a + covered value belongs (or the reverse) satisfies a row-alignment assertion by + accident. A `UInt64` covering column substituted for `_rowid` downcasts cleanly, so + the confusion is silent. `id_offset` separates the two columns. + + The first `k` rows form one tight, well-separated cluster around the query while + every other row is orders of magnitude further away. In uniformly random 128-d data + the true top-k sit at nearly identical distances, so which ones a quantizer returns + is a coin flip and a recall gate on it flakes around its threshold. Here the top-k + is unambiguous, so the recall assertion tests the search rather than the fixture. + """ + rng = np.random.default_rng(seed) + mat = rng.standard_normal((nvec, ndim)).astype(np.float32) * 10.0 + mat[:k] = mat[0] + rng.standard_normal((k, ndim)).astype(np.float32) * 0.01 + tbl = pa.table( + { + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(mat.reshape(-1), type=pa.float32()), ndim + ), + "id": pa.array([id_offset + i for i in range(nvec)], type=pa.uint64()), + "price": pa.array(rng.random(nvec) * 100), + } + ) + return tbl, mat + + +def test_create_index_covering_columns_serve_the_query_from_the_index(tmp_path): + """The payoff of the whole create-side wiring: an index built *through the Python + binding* with covering columns must answer a query for those columns out of index + storage, with values that match an independent read of the base table. + + Creation succeeding proves nothing -- nor does a correct query result on its own, + since a plain base-table take returns the same values. The plan assertions are what + separate the two, and they are deliberately both positive and negative: pinning the + exact node chain means a take inserted between the search and the output fails the + test whatever that node is called, while a bare `"LanceRead" not in plan` would go + vacuously true the moment the node is renamed -- and then the test would pass with + covering entirely dead, since values, row alignment and recall are all satisfied by + an ordinary base-table take. + """ + id_offset = 1_000_000 + tbl, mat = _covered_query_table(id_offset=id_offset) + dataset = lance.write_dataset(tbl, tmp_path) + + dataset = dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["id", "price"], + ) + desc = dataset.describe_indices()[0] + assert desc.covering_field_names == ["id", "price"] + # Covering columns never join the keyed prefix, or a consumer would believe the + # index can be searched on them. + assert desc.field_names == ["vector"] + + query = mat[0] + scanner = dataset.scanner( + nearest={"column": "vector", "q": query, "k": 10, "nprobes": 4}, + columns=["id", "price"], + with_row_id=True, + ) + plan = scanner.explain_plan(True) + # Positive: the covered plan is exactly search -> sort -> project, with nothing in + # between fetching from the base table. Asserting the whole chain (rather than only + # the absence of a node name) is what makes a rename fail loudly instead of + # silently satisfying the test. + plan_nodes = [ + line.split(":", 1)[0].strip() for line in plan.splitlines() if line.strip() + ] + assert plan_nodes == [ + "ProjectionExec", + "SortExec", + "ANNSubIndex", + "ANNIvfPartition", + ], f"unexpected node in the covered plan; plan was:\n{plan}" + # Negative, kept for the clearer failure message when the take is what came back. + assert "LanceRead" not in plan, ( + f"a covered projection must not fall back to a take; plan was:\n{plan}" + ) + + hits = scanner.to_table() + assert hits.num_rows == 10 + + # Independent read: reopen the dataset and read the base table directly, rather + # than trusting the in-memory `tbl` the index was built from. + base = lance.dataset(tmp_path).to_table(columns=["id", "price"]) + price_by_id = dict(zip(base["id"].to_pylist(), base["price"].to_pylist())) + + ids = hits["id"].to_pylist() + row_ids = hits["_rowid"].to_pylist() + prices = hits["price"].to_pylist() + # Single-fragment, step-id dataset => a correctly covered id is the row id plus the + # offset. A covered value sourced from `_rowid`, or paired with the wrong row, + # cannot satisfy this. + assert ids == [row_id + id_offset for row_id in row_ids] + assert prices == [price_by_id[i] for i in ids] + + # Serving the columns from the index is worthless if the search returns the wrong + # neighbors, so also gate on recall against brute-force ground truth. The fixture's + # planted neighborhood makes that truth unambiguous, so this gate measures the + # search and not the concentration of random high-dimensional distances. + truth = set( + (id_offset + int(i)) + for i in np.argsort(np.linalg.norm(mat - query, axis=1))[:10] + ) + recall = len(truth & set(ids)) / len(truth) + assert recall >= 0.5, f"covered recall {recall} < 0.5 (got {ids}, truth {truth})" + + +def test_create_index_uncommitted_with_covering_columns(tmp_path): + """create_index_uncommitted is a second entry point into the same PyO3 create + path. It has to forward covering_columns as well, or a distributed build silently + produces a segment with no covering columns.""" + tbl = create_table(nvec=256) + dataset = lance.write_dataset(tbl, tmp_path) + price_id = _field_id(dataset, "price") + + segment = dataset.create_index_uncommitted( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=["price"], + fragment_ids=[fragment.fragment_id for fragment in dataset.get_fragments()], + ) + assert segment.covering_fields == [price_id] + + dataset = dataset.commit_existing_index_segments("vector_idx", "vector", [segment]) + assert dataset.describe_indices()[0].segments[0].covering_fields == [price_id] + + +@pytest.mark.parametrize( + "index_type", + [pytest.param("BTREE", id="btree"), pytest.param("BITMAP", id="bitmap")], +) +def test_create_index_uncommitted_rejects_covering_columns_for_scalar( + tmp_path, index_type +): + """The segment-native scalar branch of create_index_uncommitted early-returns + before the vector params are ever built, so it cannot honour covering_columns. It + has to say so: silently returning a segment with covering_fields == [] leaves a + distributed pipeline believing it built covered scalar segments, and the covering + turns up missing only at query time.""" + dataset = lance.write_dataset(create_table(nvec=256), tmp_path) + fragment_ids = [fragment.fragment_id for fragment in dataset.get_fragments()] + + with pytest.raises(ValueError, match="only supported for vector indexes"): + dataset.create_index_uncommitted( + "price", + index_type=index_type, + fragment_ids=fragment_ids, + covering_columns=["meta"], + ) + + # The guard must not disturb the ordinary uncovered scalar path. + segment = dataset.create_index_uncommitted( + "price", index_type=index_type, fragment_ids=fragment_ids + ) + assert segment.covering_fields == [] + + +@pytest.mark.parametrize( + "covering_columns", + [ + pytest.param(["nope"], id="unknown_column"), + pytest.param(["vector"], id="indexed_column"), + pytest.param(["price", "price"], id="duplicate_column"), + ], +) +def test_create_index_covering_columns_rejected_by_core(tmp_path, covering_columns): + """Covering columns are validated in the Rust core, never in the binding. The + binding only has to let those errors reach the caller.""" + dataset = lance.write_dataset(create_table(nvec=256), tmp_path) + + with pytest.raises(ValueError, match="covering_columns"): + dataset.create_index( + "vector", + index_type="IVF_PQ", + num_partitions=4, + num_sub_vectors=16, + covering_columns=covering_columns, + ) + + def test_create_index_progress_callback_vector(tmp_path): ds = _make_sample_dataset_base(tmp_path, "vector_progress", 1500, 128) recorder = ProgressRecorder() diff --git a/python/src/dataset.rs b/python/src/dataset.rs index ffa00a78a87..0a2c448af3b 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -4968,6 +4968,7 @@ fn prepare_vector_index_params( let mut rq_params = RQBuildParams::default(); let mut index_file_version = IndexFileVersion::V3; let mut skip_transpose = false; + let mut covering_columns: Vec = Vec::new(); if let Some(kwargs) = kwargs { // Parse metric type @@ -5129,6 +5130,14 @@ fn prepare_vector_index_params( if let Some(value) = kwargs.get_item("skip_transpose")? { skip_transpose = value.extract()?; } + + // Covering ("included") columns: names of extra dataset columns stored inline in the + // index. Empty when absent. The core validates them when the index is built. + if let Some(cols) = kwargs.get_item("covering_columns")? + && !cols.is_none() + { + covering_columns = cols.extract()?; + } } let mut params = match index_type { @@ -5174,6 +5183,7 @@ fn prepare_vector_index_params( }?; params.version(index_file_version); params.skip_transpose(skip_transpose); + params.covering_columns(covering_columns); if let Some(kwargs) = kwargs && let Some(acc) = kwargs.get_item("accelerator")? { diff --git a/python/src/indices.rs b/python/src/indices.rs index 164dee21ac6..cc2f3b95fcc 100644 --- a/python/src/indices.rs +++ b/python/src/indices.rs @@ -685,6 +685,12 @@ pub struct PyIndexDescription { pub num_rows_indexed: u64, /// The details of the index pub details: PyJson, + /// The ids of the covering ("included") fields the index co-locates alongside + /// its key fields, so covered queries can skip the base-table take. Empty for + /// indexes without covering columns. All segments of one index share the set. + pub covering_fields: Vec, + /// The resolved names of `covering_fields` in the current schema. + pub covering_field_names: Vec, /// The segments of the index pub segments: Vec, /// The total size in bytes of all files across all segments @@ -705,12 +711,28 @@ impl PyIndexDescription { }) .collect(); - let segments = index + let segments: Vec = index .metadata() .iter() .map(PyIndexSegmentDescription::from_metadata) .collect(); + // All segments of one logical index declare the same covering columns + // (enforced at every commit boundary), so the first segment is authoritative. + let covering_fields = segments + .first() + .map(|segment| segment.covering_fields.clone()) + .unwrap_or_default(); + let covering_field_names = covering_fields + .iter() + .map(|field| { + dataset + .schema() + .field_path_minimal(*field) + .unwrap_or_else(|_| "".to_string()) + }) + .collect(); + let details = index.details().unwrap_or_else(|_| "{}".to_string()); Self { @@ -718,6 +740,8 @@ impl PyIndexDescription { fields: index.field_ids().to_vec(), field_names, index_type: index.index_type().to_string(), + covering_fields, + covering_field_names, segments, type_url: index.type_url().to_string(), num_rows_indexed: index.rows_indexed(), @@ -730,13 +754,19 @@ impl PyIndexDescription { #[pymethods] impl PyIndexDescription { pub fn __repr__(&self) -> String { + // `covering_fields` is printed alongside `fields` so a covered and an uncovered + // index do not render identically -- the segment repr already prints it, and + // without it here `print(ds.describe_indices()[0])` cannot answer "what does + // this index cover?" without descending into the segments. let mut repr = format!( - "IndexDescription(name='{}', type_url='{}', num_rows_indexed={}, fields={:?}, field_names={:?}, num_segments={}", + "IndexDescription(name='{}', type_url='{}', num_rows_indexed={}, fields={:?}, field_names={:?}, covering_fields={:?}, covering_field_names={:?}, num_segments={}", self.name, self.type_url, self.num_rows_indexed, self.fields, self.field_names, + self.covering_fields, + self.covering_field_names, self.segments.len() ); if let Some(byte_size) = self.total_size_bytes { @@ -764,3 +794,173 @@ pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_submodule(&indices)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StructArray}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + + /// Enough of an `IndexDescription` to drive the read-side conversion with metadata + /// no commit boundary would accept. That is the point: every write path rejects a + /// covering field id that is not a top-level field of the schema being committed + /// (`Operation::CreateIndex` in `transaction.rs`), and every schema-shrinking path + /// rejects dropping a covered column out from under a live index + /// (`reject_covered_field_subtree_change`). So `covering_field_names`'s `""` + /// fallback is only ever reached by a manifest this writer did not produce -- a + /// future format revision or a foreign writer -- and this stub is the only way to + /// hand it one. + struct StubIndexDescription { + field_ids: Vec, + segments: Vec, + } + + impl IndexDescription for StubIndexDescription { + fn name(&self) -> &str { + "stub_idx" + } + fn metadata(&self) -> &[IndexMetadata] { + &self.segments + } + fn type_url(&self) -> &str { + "" + } + fn index_type(&self) -> &str { + "Vector" + } + fn rows_indexed(&self) -> u64 { + 0 + } + fn field_ids(&self) -> &[u32] { + &self.field_ids + } + fn details(&self) -> lance_core::Result { + Ok("{}".to_string()) + } + fn total_size_bytes(&self) -> Option { + None + } + } + + fn covered_segment(fields: Vec, covering_fields: Vec) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + fields, + covering_fields, + name: "stub_idx".to_string(), + dataset_version: 1, + fragment_bitmap: None, + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + /// A dataset whose field ids deliberately do not match top-level positions: + /// `id` is field 0, the struct `nested` is field 1 (its children take 2 and 3), + /// and `price` is field 4 while sitting at top-level position 2. Resolving a + /// covered id positionally instead of by id therefore cannot accidentally + /// produce the right answer. + async fn dataset_with_offset_field_ids() -> LanceDataset { + let nested_fields = vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ]; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new( + "nested", + DataType::Struct(nested_fields.clone().into()), + false, + ), + ArrowField::new("price", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StructArray::new( + nested_fields.into(), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])) as ArrayRef, + Arc::new(Int32Array::from(vec![7, 8, 9])) as ArrayRef, + ], + None, + )), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = LanceDataset::write(reader, "memory://test", None) + .await + .unwrap(); + // Pin the ids the tests below hardcode, so a change in id assignment fails + // here rather than silently making the covered id resolve to another column. + assert_eq!(dataset.schema().field("price").unwrap().id, PRICE_FIELD_ID); + dataset + } + + /// `price`'s field id -- see [`dataset_with_offset_field_ids`]. + const PRICE_FIELD_ID: i32 = 4; + + fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + } + + #[test] + fn covering_field_names_resolve_against_the_schema() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![covered_segment( + vec![0, PRICE_FIELD_ID], + vec![PRICE_FIELD_ID], + )], + }; + + let description = PyIndexDescription::new(&index, &dataset); + + assert_eq!(description.covering_fields, vec![PRICE_FIELD_ID]); + assert_eq!(description.covering_field_names, vec!["price".to_string()]); + } + + #[test] + fn covering_field_names_fall_back_when_the_id_is_unresolvable() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + // No field carries id 99, so `field_path_minimal` returns an error. + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![covered_segment(vec![0, 99], vec![99])], + }; + + // Introspection must still succeed -- the fallback exists so a stale or + // unrecognized covering declaration cannot make `describe_indices()` error out. + let description = PyIndexDescription::new(&index, &dataset); + + assert_eq!(description.covering_fields, vec![99]); + assert_eq!( + description.covering_field_names, + vec!["".to_string()] + ); + } + + #[test] + fn covering_fields_are_empty_without_segments() { + let dataset = runtime().block_on(dataset_with_offset_field_ids()); + let index = StubIndexDescription { + field_ids: vec![0], + segments: vec![], + }; + + let description = PyIndexDescription::new(&index, &dataset); + + assert!(description.covering_fields.is_empty()); + assert!(description.covering_field_names.is_empty()); + } +}