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/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-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/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/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 151be945741..9dc843b9d6f 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -373,6 +373,41 @@ 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() +} + +/// 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"?; /// ///

@@ -410,6 +445,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 +554,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 +628,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 +706,7 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), }) } @@ -620,6 +738,7 @@ impl IvfQuantizationStorage { metadata, ivf, frag_reuse_index, + covering_schema: std::sync::OnceLock::new(), } } @@ -656,6 +775,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-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/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-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/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/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/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/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/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/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 b5803624520..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); @@ -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, @@ -265,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() @@ -277,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 @@ -289,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 }) @@ -1170,34 +1203,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 +2331,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 @@ -3637,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()); @@ -6567,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 @@ -7959,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`")] @@ -8831,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}; @@ -8953,54 +9466,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 +9577,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 +9592,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 +9800,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..ce37c91fe54 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,87 @@ 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" + ))); + } + } + match index_type { IndexType::IvfFlat => match element_type { DataType::Float16 | DataType::Float32 | DataType::Float64 => { @@ -1039,6 +1152,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 +1171,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 +1227,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 +1256,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 +1284,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 +1312,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 +1331,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 +1364,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 +1395,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 +1477,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 +1497,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 +1516,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 +1538,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 +1558,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 +1579,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 +1609,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 +1628,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 +1647,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 +1666,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 +2000,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 +2055,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 +2101,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(), @@ -2189,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}; @@ -2525,6 +2707,301 @@ 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. 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. + /// + /// 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_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::()) + .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(); + + 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" + ); + } + } + + /// 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. + /// + /// 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, + ) { + 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.with_row_id(); + 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(); + 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] async fn test_initialize_vector_index_ivf_flat() { let test_dir = TempStrDir::default(); @@ -2776,6 +3253,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..62db4a77e31 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, @@ -1544,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 @@ -1551,11 +1715,19 @@ 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). 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()); + 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, @@ -1644,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?, @@ -1696,12 +1887,15 @@ 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", )); }; - 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 +1909,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 +2003,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 +2037,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 +2057,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 +2106,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 +2298,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 +2411,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 +2489,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 +2505,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2276,6 +2516,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2286,6 +2527,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2296,6 +2538,7 @@ impl IvfIndexBuilder ivf, &row_ids, &vectors, + covering.as_ref(), new_centroids, ) .await @@ -2313,6 +2556,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 +2602,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 +3378,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..35654acafbe 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,117 @@ 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. + /// 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 +1538,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 +1786,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 +1907,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 +1923,7 @@ impl VectorIndex for IVFInd use_residual_scratch, prepared, &mut heap, + &mut covering_map, scratch, search_metrics.as_ref(), ) @@ -1730,12 +1931,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 +2127,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 +2329,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 +3157,2505 @@ 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}" + ); + } + + /// 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_covered_prefilter_results_match_non_covered() { + use arrow_array::types::UInt64Type; + use arrow_array::{Int32Array, RecordBatchIterator, StringArray}; + use arrow_buffer::NullBuffer; + + 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.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(); + + // 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(); + // 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 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" + ); + } + } + + /// 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; + + 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 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_mixed_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.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + // 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 + .create_index(&["id"], IndexType::BTree, None, &scalar, false) + .await + .unwrap(); + + 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(&format!("id < {admit_below}")).unwrap(); + scan.prefilter(true); + scan.with_row_id(); + scan.project(&["id"]).unwrap(); + + // 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" + ); + 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" + ); + } + + /// 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; + + 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]); + } + + /// 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. + #[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; + 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. + /// + /// 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 batches = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(batches, test_uri, None).await.unwrap(); + + 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( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + + // 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.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() > 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 + + /// `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). + #[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_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; + + 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(); + + // 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!( + !plan.contains("LanceRead"), + "covered projection ['id'] should skip the take after optimize; plan:\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::(); + + // 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 + /// 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()); + } + + /// 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 + /// 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 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) + /// 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"); + } + } + + /// 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 + /// 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/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 7a334198610..75f97177de0 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,61 @@ 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(); + // 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 { + // 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 +1500,8 @@ impl ANNIvfSubIndexExec { prefilter_source, overlay_block: None, external_mask: None, + output_schema, + has_covered, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1513,6 +1581,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,20 +1602,26 @@ 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()), } } - 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) { @@ -1763,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) }) @@ -1782,6 +1866,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 +1911,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. @@ -2022,13 +2130,146 @@ impl ANNIvfSubIndexExec { seg_mask, ) .await?; - state.record_batch(&batch); + state.record_batch(&batch)?; Ok(batch) } }) .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()`. + // 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() + .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, projection)? + .include_row_id() + .execute() + .await?; + + // 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)) + } } impl ExecutionPlan for ANNIvfSubIndexExec { @@ -2037,7 +2278,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 +2325,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 +2345,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 +2449,25 @@ 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`). + // 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(); + let recon_state = state.clone(); + 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, per_index_stream @@ -2277,6 +2543,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state, target_partitions, seg_mask, + has_covered, ); DataFusionResult::Ok(early_search.chain(late_search).boxed()) } @@ -2285,6 +2552,66 @@ 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). + 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 +2687,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 +2713,7 @@ impl MultivectorScoringExec { Ok(Self { inputs, query, + output_schema, properties, }) } @@ -2399,7 +2738,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 +2773,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 +2798,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 +2815,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 +2840,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 +2887,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 +3004,147 @@ 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(), + ); + + // 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![0, 9999], + 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], + }; + index + .validate_covering_fields() + .expect("the fixture must be a legal covering declaration"); + + 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}" + ); + } + + /// 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(); @@ -2826,7 +3371,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 +3744,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)) } @@ -3384,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( @@ -3407,6 +4022,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>() .await @@ -3418,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() @@ -3522,6 +4164,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3601,6 +4244,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + false, ) .try_collect::>() .await @@ -3614,6 +4258,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 +4483,7 @@ mod tests { state.clone(), usize::MAX, None, + false, ) .try_collect::>(); @@ -3653,6 +4503,7 @@ mod tests { state, usize::MAX, None, + false, ) .try_collect::>(); @@ -4118,6 +4969,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