diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index c6f448040c2..c96a3db915f 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -8,7 +8,7 @@ use crate::dataset::ROW_ID; use crate::dataset::WriteDestination; use crate::dataset::optimize::{CompactionOptions, compact_files}; use crate::dataset::transaction::{DataReplacementGroup, Operation}; -use crate::dataset::{AutoCleanupParams, MergeInsertBuilder, ProjectionRequest}; +use crate::dataset::{AutoCleanupParams, MergeInsertBuilder, ProjectionRequest, UpdateBuilder}; use crate::index::DatasetIndexExt; use crate::{Dataset, Error}; use lance_core::ROW_ADDR; @@ -28,7 +28,7 @@ use arrow_array::{ ArrayRef, Float32Array, Int32Array, ListArray, RecordBatchIterator, StringArray, types::Int32Type, }; -use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; use lance_arrow::BLOB_META_KEY; use lance_core::utils::tempfile::{TempDir, TempStrDir}; use lance_datafusion::utils::reader_to_stream; @@ -1627,6 +1627,121 @@ async fn test_merge_insert_with_reordered_columns_and_index() { final_dataset.validate().await.unwrap(); } +/// With stable row ids, updating a top-level struct column keeps a scalar index on a +/// nested child field correct. The update API rejects nested column references, so a +/// nested field can only be changed by setting its whole struct column; that update must +/// not wrongly extend the child-field index over the rewritten fragment (which would +/// leave the updated value unscanned and silently dropped). +#[tokio::test] +async fn test_update_struct_column_keeps_nested_index() { + let struct_fields = Fields::from(vec![ArrowField::new("x", DataType::Int32, true)]); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("s", DataType::Struct(struct_fields.clone()), true), + ])); + let s_arr = StructArray::new( + struct_fields.clone(), + vec![Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef], + None, + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(s_arr) as ArrayRef, + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write( + reader, + "memory://test_update_nested_index", + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + // BTree index on the NESTED field `s.x`. + dataset + .create_index( + &["s.x"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let pre = dataset + .scan() + .filter("s.x = 20") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(pre.num_rows(), 1, "precondition: s.x=20 should match id=2"); + + // Nested column references are rejected by `set`, so update the whole struct column + // `s` for id=2, changing s.x 20 -> 999. + let update_result = UpdateBuilder::new(Arc::new(dataset.clone())) + .update_where("id = 2") + .unwrap() + .set("s", "named_struct('x', cast(999 as int))") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + let dataset = update_result.new_dataset; + + // The nested `s.x` index must NOT be extended to the rewritten fragment: its + // effective coverage stays {0}, so the rewritten fragment is left unindexed and + // fully scanned. + let sx_idx = dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|i| i.fields.len() == 1) + .expect("nested s.x index") + .clone(); + let effective = sx_idx + .effective_fragment_bitmap(&dataset.fragment_bitmap) + .expect("index has a fragment bitmap"); + assert_eq!( + effective.iter().collect::>(), + vec![0], + "nested-field index must not be extended to the rewritten fragment" + ); + + // The updated value must be found, and the stale value gone. + let new = dataset + .scan() + .filter("s.x = 999") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + new.num_rows(), + 1, + "updated value s.x=999 must be found after the struct-column update" + ); + let old = dataset + .scan() + .filter("s.x = 20") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(old.num_rows(), 0, "s.x=20 should no longer match any row"); +} + /// DataReplacement should invalidate index fragment bitmaps for replaced fields. #[tokio::test] async fn test_data_replacement_invalidates_index_bitmap() { diff --git a/rust/lance/src/dataset/write/update.rs b/rust/lance/src/dataset/write/update.rs index 6672c58db1f..f8e71b834a7 100644 --- a/rust/lance/src/dataset/write/update.rs +++ b/rust/lance/src/dataset/write/update.rs @@ -34,6 +34,16 @@ use lance_table::format::{Fragment, RowIdMeta}; use roaring::RoaringTreemap; use snafu::ResultExt; +/// Collect a field id and all of its descendant field ids (pre-order). A struct +/// column update rewrites the whole subtree, so an index on any descendant must be +/// treated as modified. +fn collect_subtree_field_ids(field: &lance_core::datatypes::Field, out: &mut Vec) { + out.push(field.id as u32); + for child in &field.children { + collect_subtree_field_ids(child, out); + } +} + /// Build an update operation. /// /// This operation is similar to SQL's UPDATE statement. It allows you to change @@ -386,10 +396,14 @@ impl UpdateJob { dataset: Arc, update_data: UpdateData, ) -> Result { + // Updated columns are top-level (nested references are rejected by `set`), but a + // struct-column update rewrites all of its descendants. Collect the full field + // subtree so an index on a nested child field is recognized as modified and not + // wrongly extended over the rewritten fragment. let mut fields_for_preserving_frag_bitmap = Vec::new(); for column_name in self.updates.keys() { - if let Ok(field_id) = dataset.schema().field_id(column_name) { - fields_for_preserving_frag_bitmap.push(field_id as u32); + if let Some(field) = dataset.schema().field(column_name) { + collect_subtree_field_ids(field, &mut fields_for_preserving_frag_bitmap); } }