From 5a4f387101ea19c161778426fcfeadb43c51b4d8 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 01:15:11 -0400 Subject: [PATCH 01/48] fix(mem_wal): replay a WAL entry whose schema has since moved MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A replayed entry was re-labelled to the current storage schema by position: the entry's columns were taken in order and rebound to whatever the schema now declares. That only holds while the schema is the one the entry was written under. It is not, after a column is added or dropped. An entry that predates an added column is one column short, and the rebind fails on the width — the shard cannot be opened at all. An entry that predates a dropped column is one column long, and the rebind silently stores each remaining column under its neighbour's name, which is worse: the open succeeds and the rows are wrong. Match by name instead. A column the schema declares and the entry does not carry is filled with nulls, which is what a column added later means for rows written before it; non-primary-key columns are nullable in the storage schema whatever the base table declares, so the null is always representable. A column the entry carries and the schema no longer declares is dropped. Two things stay errors, because no fill is right for them: a missing primary key, which cannot be invented, and a column whose type changed. Live writes are unaffected. They are checked against the logical schema before they reach here, so every column is present and this still only appends `_tombstone`. --- rust/lance/src/dataset/mem_wal/write.rs | 180 +++++++++++++++++++++--- 1 file changed, 164 insertions(+), 16 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index d0d1246ea83..ac451e00c89 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -1389,6 +1389,9 @@ async fn replay_memtable_from_wal( manifest: &ShardManifest, base_generation: u64, mut make_memtable: impl FnMut(u64, usize) -> Result, + // Conforming a replayed entry needs these: a primary key the entry does not + // carry cannot be filled with a null. + pk_columns: &[String], flusher: &MemTableFlusher, wal_flusher: &WalFlusher, index_configs: &[MemIndexConfig], @@ -1432,7 +1435,7 @@ async fn replay_memtable_from_wal( let batches = entry .batches .into_iter() - .map(|b| ensure_tombstone_column(b, &storage_schema)) + .map(|b| conform_to_storage_schema(b, &storage_schema, pk_columns)) .collect::>>()?; // Seal + flush on the same criteria the live path uses, measured @@ -1598,24 +1601,50 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// Re-label `batch` to the storage schema, injecting `_tombstone = false` when -/// absent — callers pass logical-shaped batches, and WAL entries written before -/// deletes existed lack the column. +/// Re-label `batch` to the storage schema, matching columns by **name**. +/// +/// A column the schema declares and the batch does not carry is filled with +/// typed nulls; `_tombstone` is filled with `false`. A column the batch carries +/// and the schema does not declare is dropped. +/// +/// Both cases are what a replayed WAL entry looks like after the table's schema +/// moved: an entry predates a column added since, and carries one dropped +/// since. Matching by position instead would reject the first outright and +/// store the second under its neighbour's name. Live writes reach here already +/// checked against the logical schema, so for them every column is present and +/// this only appends `_tombstone`. /// -/// A batch that already carries `_tombstone` is re-labeled too, so an entry -/// written under an older storage schema replays into the current one. -fn ensure_tombstone_column( +/// A primary key the batch does not carry is an error — there is no value to +/// invent — and so is a column whose type does not match the schema. +fn conform_to_storage_schema( batch: RecordBatch, storage_schema: &Arc, + pk_columns: &[String], ) -> Result { let n = batch.num_rows(); - let mut columns: Vec = batch.columns().to_vec(); - if batch.schema().column_with_name(TOMBSTONE).is_none() { - columns.push(Arc::new(BooleanArray::from(vec![false; n]))); + let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); + for field in storage_schema.fields() { + let name = field.name(); + if let Some(column) = batch.column_by_name(name) { + columns.push(column.clone()); + } else if name == TOMBSTONE { + columns.push(Arc::new(BooleanArray::from(vec![false; n]))); + } else if pk_columns.iter().any(|c| c == name) { + return Err(Error::invalid_input(format!( + "batch is missing primary key column '{}' declared by the storage schema", + name + ))); + } else { + // Non-primary-key columns are nullable in the storage schema + // whatever the base table declares (`relax_non_pk_nullability`), so + // a null stands in for a value the entry never held. + columns.push(new_null_array(field.data_type(), n)); + } } RecordBatch::try_new(storage_schema.clone(), columns).map_err(|e| { Error::invalid_input(format!( - "failed to inject _tombstone column (does the batch match the base table schema?): {}", + "failed to conform a batch to the storage schema \ + (does the batch match the base table schema?): {}", e )) }) @@ -2371,6 +2400,7 @@ impl ShardWriter { manifest, manifest.current_generation, make_bound_memtable, + &pk_columns, &flusher, &wal_flusher, index_configs, @@ -2639,7 +2669,9 @@ impl ShardWriter { // `_tombstone`. let batches = batches .into_iter() - .map(|b| ensure_tombstone_column(b, &writer_state.schema)) + .map(|b| { + conform_to_storage_schema(b, &writer_state.schema, &writer_state.pk_columns) + }) .collect::>>()?; self.put_memtable(batches, state, writer_state, backpressure) .await @@ -2764,7 +2796,9 @@ impl ShardWriter { // Mirrors `put`. let batches = batches .into_iter() - .map(|b| ensure_tombstone_column(b, &writer_state.schema)) + .map(|b| { + conform_to_storage_schema(b, &writer_state.schema, &writer_state.pk_columns) + }) .collect::>>()?; self.put_memtable_no_wait(batches, state, writer_state, backpressure) .await @@ -4651,10 +4685,11 @@ mod tests { } #[test] - fn test_ensure_tombstone_column_injects_false() { + fn test_conform_injects_tombstone_false() { let base = create_test_schema(); let storage = schema_with_tombstone(&base); - let out = ensure_tombstone_column(create_test_batch(&base, 0, 3), &storage).unwrap(); + let pk = ["id".to_string()]; + let out = conform_to_storage_schema(create_test_batch(&base, 0, 3), &storage, &pk).unwrap(); assert_eq!(out.schema(), storage); let ts = out .column_by_name(TOMBSTONE) @@ -4667,10 +4702,123 @@ mod tests { "put injects _tombstone = false" ); // Idempotent: a batch already carrying the column passes through. - let again = ensure_tombstone_column(out.clone(), &storage).unwrap(); + let again = conform_to_storage_schema(out.clone(), &storage, &pk).unwrap(); assert_eq!(again.schema(), out.schema()); } + /// A WAL entry written before a column was added still replays: the column + /// it never held becomes null rather than a width mismatch. + #[test] + fn test_conform_fills_a_column_added_since_the_entry() { + let pk = ["id".to_string()]; + let entry = create_test_batch(&create_test_schema(), 0, 2); + + let widened = ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("added_later", DataType::Int64, true), + ]); + let storage = schema_with_tombstone(&widened); + + let out = conform_to_storage_schema(entry, &storage, &pk).unwrap(); + assert_eq!(out.schema(), storage); + assert_eq!(out.num_rows(), 2); + let added = out.column_by_name("added_later").unwrap(); + assert_eq!(added.null_count(), 2, "the new column replays as all-null"); + let ids = out + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[0, 1], "the entry's own columns survive"); + } + + /// A WAL entry written before a column was dropped still replays: the + /// column the schema no longer declares is left behind, and the columns + /// that remain keep their own values rather than their neighbour's. + #[test] + fn test_conform_drops_a_column_removed_since_the_entry() { + let pk = ["id".to_string()]; + let wide = ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("dropped_later", DataType::Utf8, true), + Field::new("name", DataType::Utf8, true), + ]); + let entry = RecordBatch::try_new( + Arc::new(wide), + vec![ + Arc::new(Int32Array::from(vec![7, 8])), + Arc::new(StringArray::from(vec!["gone", "gone"])), + Arc::new(StringArray::from(vec!["kept-7", "kept-8"])), + ], + ) + .unwrap(); + + let storage = schema_with_tombstone(&create_test_schema()); + let out = conform_to_storage_schema(entry, &storage, &pk).unwrap(); + assert_eq!(out.schema(), storage); + let names = out + .column_by_name("name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + (names.value(0), names.value(1)), + ("kept-7", "kept-8"), + "positional re-labelling would have stored `dropped_later` as `name`" + ); + } + + /// A primary key is the one column a null cannot stand in for. + #[test] + fn test_conform_refuses_an_entry_missing_a_primary_key() { + let storage = schema_with_tombstone(&create_test_schema()); + let keyless = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + "name", + DataType::Utf8, + true, + )])), + vec![Arc::new(StringArray::from(vec!["a"]))], + ) + .unwrap(); + + let error = conform_to_storage_schema(keyless, &storage, &["id".to_string()]).unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected InvalidInput, got {error:?}" + ); + assert!( + error.to_string().contains("id"), + "the error should name the missing key: {error}" + ); + } + + /// A column whose type changed is a conflict, not a drift to paper over. + #[test] + fn test_conform_refuses_a_retyped_column() { + let storage = schema_with_tombstone(&create_test_schema()); + let retyped = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Boolean, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(BooleanArray::from(vec![true])), + ], + ) + .unwrap(); + + let error = conform_to_storage_schema(retyped, &storage, &["id".to_string()]).unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected InvalidInput, got {error:?}" + ); + } + #[test] fn test_build_tombstone_batch_shape() { let storage = schema_with_tombstone(&create_test_schema()); From 3d8965059874791f83df35607319c8e6d2511cd5 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 16:01:37 -0400 Subject: [PATCH 02/48] fix(mem_wal): carry field ids into the memtable schema, so a rename keeps its rows Replay matches a WAL entry to the storage schema by name, which holds for a column added or dropped but not for one renamed. A rename is the one change that moves a column's name while keeping its identity: the name is mutated in place and the field id is untouched, which is why data files need no rewrite -- they list field ids, not names. Matching on name alone puts the fresh tier at odds with that. The entry carries the old name, the schema declares the new one, and the column is nulled under its new name while the old one is dropped: every row still in the memtable loses the value. Carry the id. `From<&Field> for ArrowField` drops it, so the memtable's storage schema stamps it into field metadata itself, and Arrow IPC preserves field metadata, so every WAL entry written under that schema carries it too. Conform then matches on id where both sides have one and falls back to name where they do not -- entries written before this have no ids, so both tiers are needed. Scoped to the memtable path on purpose. Emitting the id from the global Arrow conversion would change every schema Lance hands out, including for callers that compare schemas for equality; this changes only the schema the memtable and its WAL entries are written under. --- rust/lance-core/src/datatypes.rs | 2 +- rust/lance/src/dataset/mem_wal.rs | 43 +++++++++++++++++++++++++ rust/lance/src/dataset/mem_wal/api.rs | 2 +- rust/lance/src/dataset/mem_wal/write.rs | 43 ++++++++++++++++++++++--- 4 files changed, 84 insertions(+), 6 deletions(-) diff --git a/rust/lance-core/src/datatypes.rs b/rust/lance-core/src/datatypes.rs index dcad595e5aa..d86b764b305 100644 --- a/rust/lance-core/src/datatypes.rs +++ b/rust/lance-core/src/datatypes.rs @@ -18,7 +18,7 @@ mod schema; use crate::{Error, Result}; pub use field::{ - BlobVersion, Encoding, Field, LANCE_UNENFORCED_CLUSTERING_KEY_POSITION, + BlobVersion, Encoding, Field, LANCE_FIELD_ID_KEY, LANCE_UNENFORCED_CLUSTERING_KEY_POSITION, LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION, NullabilityComparison, OnTypeMismatch, SchemaCompareOptions, }; diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 79c7fac8272..cf4f856a2b0 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -48,6 +48,10 @@ pub mod write; use std::sync::Arc; +use std::collections::HashMap; + +use lance_core::datatypes::{LANCE_FIELD_ID_KEY, Schema}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; /// Column name for the mem_wal tombstone (delete sentinel) marker. @@ -119,6 +123,45 @@ pub fn relax_non_pk_nullability( /// Idempotent: a schema that already carries `_tombstone` (a reopen/replay /// path) is returned unchanged. Schema-level metadata and per-field metadata /// (e.g. the `lance-schema:unenforced-primary-key` marker) are preserved. +/// The schema's Arrow form, with each field's id carried in its metadata. +/// +/// `From<&Field> for ArrowField` drops the id, which leaves everything +/// downstream matching on name alone. That holds until a column is renamed: the +/// name is the part a rename changes and the id is the part it keeps, so a +/// name-only match loses the column. Data files have always been addressed by +/// id (`DataFile.fields` lists them); carrying it into the memtable's storage +/// schema puts the fresh tier on the same footing, and Arrow IPC keeps field +/// metadata, so every WAL entry written under this schema carries it too. +/// +/// Scoped to the memtable path deliberately: emitting the id from the global +/// Arrow conversion would change every schema Lance hands out, including for +/// callers that compare schemas for equality. +pub(crate) fn arrow_schema_with_field_ids(schema: &Schema) -> ArrowSchema { + let arrow: ArrowSchema = schema.into(); + let ids: HashMap<&str, i32> = schema + .fields + .iter() + .map(|f| (f.name.as_str(), f.id)) + .collect(); + let fields: Vec = arrow + .fields() + .iter() + .map(|field| { + let Some(id) = ids + .get(field.name().as_str()) + .copied() + .filter(|id| *id >= 0) + else { + return field.as_ref().clone(); + }; + let mut metadata = field.metadata().clone(); + metadata.insert(LANCE_FIELD_ID_KEY.to_string(), id.to_string()); + field.as_ref().clone().with_metadata(metadata) + }) + .collect(); + ArrowSchema::new_with_metadata(fields, arrow.metadata().clone()) +} + pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { if base.column_with_name(TOMBSTONE).is_some() { return Arc::new(base.clone()); diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index ef55d20d5ab..bceedb2919a 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -702,7 +702,7 @@ impl DatasetMemWalExt for Dataset { base_path, base_uri, config, - Arc::new(self.schema().into()), + Arc::new(super::arrow_schema_with_field_ids(self.schema())), index_configs, ) .await diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index ac451e00c89..39ec5e3440b 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -20,9 +20,9 @@ use std::time::{Duration, Instant}; use arc_swap::ArcSwap; use arrow_array::{ArrayRef, BooleanArray, RecordBatch, new_null_array}; -use arrow_schema::Schema as ArrowSchema; +use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use async_trait::async_trait; -use lance_core::datatypes::Schema; +use lance_core::datatypes::{LANCE_FIELD_ID_KEY, Schema}; use lance_core::{Error, Result}; use lance_index::mem_wal::ShardManifest; use lance_index::vector::hnsw::builder::HnswBuildParams; @@ -1601,12 +1601,33 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// Re-label `batch` to the storage schema, matching columns by **name**. +/// The lance field id an Arrow field carries, if it carries one. +/// +/// Lance already reads this key when converting Arrow to its own schema; the +/// memtable's storage schema carries it so that a WAL entry, which is Arrow IPC +/// and so keeps field metadata, stays addressable by id rather than by name +/// alone. Base data files have always been addressed this way +/// (`DataFile.fields` is a list of ids); this brings the fresh tier alongside. +fn field_id_of(field: &ArrowField) -> Option { + field + .metadata() + .get(LANCE_FIELD_ID_KEY) + .and_then(|v| v.parse::().ok()) + .filter(|id| *id >= 0) +} + +/// Re-label `batch` to the storage schema, matching columns by **field id** +/// where both sides carry one, and by **name** otherwise. /// /// A column the schema declares and the batch does not carry is filled with /// typed nulls; `_tombstone` is filled with `false`. A column the batch carries /// and the schema does not declare is dropped. /// +/// Ids are tried first because a rename keeps the id and changes the name, so a +/// name match would null the new name and drop the old one — the column's values +/// lost. Entries written before ids were carried have none, and fall back to the +/// name match, which is why both tiers exist. +/// /// Both cases are what a replayed WAL entry looks like after the table's schema /// moved: an entry predates a column added since, and carries one dropped /// since. Matching by position instead would reject the first outright and @@ -1622,10 +1643,24 @@ fn conform_to_storage_schema( pk_columns: &[String], ) -> Result { let n = batch.num_rows(); + // A field id survives a rename; a name does not. Matching on it first is + // what keeps a renamed column's values attached to the column, instead of + // nulling the new name and dropping the old one. Absent on entries written + // before ids were carried, which is why the name match below remains. + let by_field_id: HashMap = batch + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(i, f)| field_id_of(f).map(|id| (id, batch.column(i)))) + .collect(); + let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); for field in storage_schema.fields() { let name = field.name(); - if let Some(column) = batch.column_by_name(name) { + if let Some(column) = field_id_of(field).and_then(|id| by_field_id.get(&id)) { + columns.push((*column).clone()); + } else if let Some(column) = batch.column_by_name(name) { columns.push(column.clone()); } else if name == TOMBSTONE { columns.push(Arc::new(BooleanArray::from(vec![false; n]))); From effa6dcbd94943640eac6b0935d7f4bd5372adfe Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 17:00:12 -0400 Subject: [PATCH 03/48] fix(mem_wal): cast a replayed column whose type changed `alter_columns` casts the rows already in the base table rather than rewriting them, so a column's type can move under a shard that still holds entries written at the old type. Conform rejected those outright, which left the shard unopenable: the entry is not wrong, it is just older than the schema. Cast it, under the same `safe: false` option `alter_columns` uses, so a replayed row lands in the state it would have had if written after the change and a lossy cast stays an error rather than a column of nulls. Both halves of the table then agree on what the change meant. --- rust/lance/src/dataset/mem_wal/write.rs | 144 +++++++++++++++++++++--- 1 file changed, 130 insertions(+), 14 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 39ec5e3440b..5b7937e48be 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -19,7 +19,9 @@ use std::sync::{Arc, RwLock as StdRwLock}; use std::time::{Duration, Instant}; use arc_swap::ArcSwap; +use arrow::compute::CastOptions; use arrow_array::{ArrayRef, BooleanArray, RecordBatch, new_null_array}; +use arrow_cast::cast_with_options; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use async_trait::async_trait; use lance_core::datatypes::{LANCE_FIELD_ID_KEY, Schema}; @@ -1635,8 +1637,13 @@ fn field_id_of(field: &ArrowField) -> Option { /// checked against the logical schema, so for them every column is present and /// this only appends `_tombstone`. /// -/// A primary key the batch does not carry is an error — there is no value to -/// invent — and so is a column whose type does not match the schema. +/// A column whose type changed is cast, the same way `alter_columns` casts the +/// base data, so a replayed row lands in the state it would have had if it had +/// been written after the change. A cast that would lose information is an +/// error, not a silent null. +/// +/// A primary key the batch does not carry stays an error — there is no value to +/// invent. fn conform_to_storage_schema( batch: RecordBatch, storage_schema: &Arc, @@ -1658,10 +1665,33 @@ fn conform_to_storage_schema( let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); for field in storage_schema.fields() { let name = field.name(); - if let Some(column) = field_id_of(field).and_then(|id| by_field_id.get(&id)) { - columns.push((*column).clone()); - } else if let Some(column) = batch.column_by_name(name) { - columns.push(column.clone()); + let carried = field_id_of(field) + .and_then(|id| by_field_id.get(&id).map(|c| (*c).clone())) + .or_else(|| batch.column_by_name(name).cloned()); + if let Some(column) = carried { + columns.push(if column.data_type() == field.data_type() { + column + } else { + // `safe: false` so a lossy cast is an error rather than a + // column of nulls -- the same option `alter_columns` casts the + // base data under, so both halves of the table agree. + cast_with_options( + &column, + field.data_type(), + &CastOptions { + safe: false, + ..Default::default() + }, + ) + .map_err(|e| { + Error::invalid_input(format!( + "column '{name}' was written as {} and the schema now declares {}, \ + which it cannot be cast to: {e}", + column.data_type(), + field.data_type(), + )) + })? + }); } else if name == TOMBSTONE { columns.push(Arc::new(BooleanArray::from(vec![false; n]))); } else if pk_columns.iter().any(|c| c == name) { @@ -4598,7 +4628,7 @@ pub fn new_shared_stats() -> SharedWriteStats { mod tests { use super::*; use crate::dataset::mem_wal::test_util::failing_memory_store; - use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, StringArray}; + use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, Int64Array, StringArray}; use arrow_schema::{DataType, Field}; use lance_core::FenceReason; use rstest::rstest; @@ -4831,27 +4861,113 @@ mod tests { ); } - /// A column whose type changed is a conflict, not a drift to paper over. + /// A widened type is cast, matching what `alter_columns` did to the rows + /// already in the base table. #[test] - fn test_conform_refuses_a_retyped_column() { - let storage = schema_with_tombstone(&create_test_schema()); - let retyped = RecordBatch::try_new( + fn test_conform_casts_a_widened_column() { + let widened = ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("count", DataType::Int64, true), + ]); + let storage = schema_with_tombstone(&widened); + + let narrow = RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![ Field::new("id", DataType::Int32, false), - Field::new("name", DataType::Boolean, true), + Field::new("name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), ])), vec![ Arc::new(Int32Array::from(vec![1])), - Arc::new(BooleanArray::from(vec![true])), + Arc::new(StringArray::from(vec!["a"])), + Arc::new(Int32Array::from(vec![7])), ], ) .unwrap(); - let error = conform_to_storage_schema(retyped, &storage, &["id".to_string()]).unwrap_err(); + let out = conform_to_storage_schema(narrow, &storage, &["id".to_string()]).unwrap(); + assert_eq!(out.schema(), storage); + let counts = out + .column_by_name("count") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(counts.value(0), 7, "the value survives the widening"); + } + + /// A cast that would lose the value is an error, not a column of nulls. + #[test] + fn test_conform_refuses_a_lossy_retype() { + let numeric = schema_with_tombstone(&ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Int32, true), + ])); + let textual = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec!["not a number"])), + ], + ) + .unwrap(); + + let error = conform_to_storage_schema(textual, &numeric, &["id".to_string()]).unwrap_err(); assert!( matches!(error, Error::InvalidInput { .. }), "expected InvalidInput, got {error:?}" ); + assert!( + error.to_string().contains("name"), + "the error should name the column: {error}" + ); + } + + /// A field id outlives a rename, so matching on it keeps the column's rows + /// where matching on the name would null them. + #[test] + fn test_conform_follows_a_field_id_through_a_rename() { + fn with_id(field: ArrowField, id: i32) -> ArrowField { + let mut metadata = field.metadata().clone(); + metadata.insert(LANCE_FIELD_ID_KEY.to_string(), id.to_string()); + field.with_metadata(metadata) + } + + // The entry was written while field 1 was called `before`. + let entry = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + with_id(Field::new("id", DataType::Int32, false), 0), + with_id(Field::new("before", DataType::Utf8, true), 1), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec!["kept"])), + ], + ) + .unwrap(); + + // The schema now calls field 1 `after`. + let storage = schema_with_tombstone(&ArrowSchema::new(vec![ + with_id(Field::new("id", DataType::Int32, false), 0), + with_id(Field::new("after", DataType::Utf8, true), 1), + ])); + + let out = conform_to_storage_schema(entry, &storage, &["id".to_string()]).unwrap(); + let after = out + .column_by_name("after") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + after.value(0), + "kept", + "the id should carry the value to the new name; a name match would null it" + ); } #[test] From 26c5004b9837797a77e80727fe9e3f6ac984693c Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 21:32:34 -0400 Subject: [PATCH 04/48] test(mem_wal): cover a struct column matched by its own field id Ids are carried for top-level fields, which is the granularity conform works at. A struct column is taken or it is not, and a change inside it is a change to the column's type. --- rust/lance/src/dataset/mem_wal/write.rs | 45 +++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 5b7937e48be..a1b11b774b2 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -4927,6 +4927,51 @@ mod tests { ); } + /// A struct column is matched whole, by the id on the column itself. + /// + /// Ids are carried for top-level fields, which is the granularity conform + /// works at: a column is taken or it is not. A change inside the struct + /// changes the column's type, and is handled as a type change. + #[test] + fn test_conform_matches_a_struct_column_by_its_own_id() { + fn with_id(field: ArrowField, id: i32) -> ArrowField { + let mut metadata = field.metadata().clone(); + metadata.insert(LANCE_FIELD_ID_KEY.to_string(), id.to_string()); + field.with_metadata(metadata) + } + + let child = Arc::new(ArrowField::new("inner", DataType::Int32, true)); + let struct_type = DataType::Struct(vec![child.clone()].into()); + let values: ArrayRef = Arc::new(arrow_array::StructArray::from(vec![( + child, + Arc::new(Int32Array::from(vec![42])) as ArrayRef, + )])); + + let entry = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + with_id(Field::new("id", DataType::Int32, false), 0), + with_id(Field::new("before", struct_type.clone(), true), 1), + ])), + vec![Arc::new(Int32Array::from(vec![1])), values], + ) + .unwrap(); + + // The column was renamed; its type, and so its children, are unchanged. + let storage = schema_with_tombstone(&ArrowSchema::new(vec![ + with_id(Field::new("id", DataType::Int32, false), 0), + with_id(Field::new("after", struct_type, true), 1), + ])); + + let out = conform_to_storage_schema(entry, &storage, &["id".to_string()]).unwrap(); + let after = out.column_by_name("after").expect("renamed struct column"); + assert_eq!( + after.null_count(), + 0, + "the struct column should carry values" + ); + assert_eq!(out.schema(), storage); + } + /// A field id outlives a rename, so matching on it keeps the column's rows /// where matching on the name would null them. #[test] From 176827b9190a7064d108ecf85073fab03c2ce69b Mon Sep 17 00:00:00 2001 From: XYZhan Date: Fri, 11 Sep 2026 23:54:47 -0400 Subject: [PATCH 05/48] docs(mem_wal): say what a schema without field ids costs a caller `ShardWriter::open` takes whatever schema it is given. One built without field ids still works -- a replayed entry falls back to matching by name -- but a column renamed while that shard holds rows then reads as null for every one of them, and nothing about the call says so. --- rust/lance/src/dataset/mem_wal/write.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index a1b11b774b2..949eea193f5 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -2150,6 +2150,12 @@ impl ShardWriter { /// /// The `base_path` should come from `ObjectStore::from_uri()` to ensure /// WAL files are written inside the dataset directory. + /// + /// `schema` carrying each field's id under [`LANCE_FIELD_ID_KEY`] in its + /// field metadata is what lets a replayed entry be matched to a column that + /// has since been renamed. Without them, a replayed entry is matched by + /// name, and a renamed column reads as null for every row the memtable + /// still holds. #[instrument(name = "sw_open", level = "info", skip_all, fields(shard_id = %config.shard_id, index_count = index_configs.len()))] pub async fn open( object_store: Arc, From 6f2256f224ba0c9a47c546a626b7a340970d1474 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 14:47:14 -0400 Subject: [PATCH 06/48] fix(mem_wal): open a shard whose maintained index is gone A maintained set is fixed when the write spec is installed and cannot be edited afterwards, so an index that disappears -- dropped outright, or carried away with the column it covered -- is named by a set that can never stop naming it. Refusing to build the configs then refuses the claim, and a table whose claim cannot be built serves no reads at all, over a condition that costs the fresh tier one index. Skip it and say so. The base index is gone for every reader; the fresh tier simply has nothing left to keep in step with. Validating a set before it is installed still rejects an unresolvable name. That is the operator's mistake and the one moment it can still be corrected. --- rust/lance/src/dataset/mem_wal/api.rs | 59 +++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index bceedb2919a..ae39bd30a07 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -678,8 +678,13 @@ impl DatasetMemWalExt for Dataset { // Get maintained_indexes from the MemWalIndex details let maintained_indexes = &mem_wal_index.details.maintained_indexes; - let index_configs = - build_index_configs(self, maintained_indexes, &config.hnsw_params).await?; + let index_configs = build_index_configs( + self, + maintained_indexes, + &config.hnsw_params, + OnMissingIndex::Skip, + ) + .await?; // Set shard_id in config config.shard_id = shard_id; @@ -714,10 +719,24 @@ impl DatasetMemWalExt for Dataset { /// Shared by [`DatasetMemWalExt::mem_wal_writer`] and /// [`validate_maintained_indexes`], so a set that validates is one the writer /// can build. +/// Whether an index the set names but the dataset does not have is fatal. +#[derive(Clone, Copy, PartialEq, Eq)] +enum OnMissingIndex { + /// Validating a set before it is installed: a name that resolves to nothing + /// is the operator's mistake, and the only moment it can still be corrected. + Reject, + /// Opening a shard against a set installed earlier: the index may have been + /// dropped since, or carried away with the column it covered. The set + /// cannot be edited, so rejecting here refuses every read on the table for + /// something that costs only the fresh tier's copy of one index. + Skip, +} + async fn build_index_configs( dataset: &Dataset, index_names: &[String], hnsw_params: &HashMap, + on_missing: OnMissingIndex, ) -> Result> { let mut index_configs = Vec::with_capacity(index_names.len()); for index_name in index_names { @@ -729,13 +748,31 @@ async fn build_index_configs( .load_indices_by_name(index_name) .await? .into_iter() - .next() - .ok_or_else(|| { - Error::invalid_input(format!( + .next(); + + // An index the maintained set names and the dataset no longer has: + // dropped outright, or carried away with the column it covered. The set + // is fixed when the write spec is installed and cannot be edited + // afterwards, so refusing here refuses the claim -- and a table whose + // claim cannot be built serves no reads at all, for a condition that + // costs only the fresh tier's copy of one index. + // + // Serve without it instead. The base index is gone for everyone; the + // fresh tier simply has nothing to keep in step with. + let Some(index_meta) = index_meta else { + if on_missing == OnMissingIndex::Reject { + return Err(Error::invalid_input(format!( "Index '{}' from maintained_indexes not found on dataset", index_name - )) - })?; + ))); + } + log::warn!( + "index '{}' is named by maintained_indexes but is not on the dataset; \ + the fresh tier will not maintain it", + index_name + ); + continue; + }; // Detect index kind and create appropriate config let type_url = index_meta @@ -782,7 +819,13 @@ async fn build_index_configs( pub async fn validate_maintained_indexes(dataset: &Dataset, index_names: &[String]) -> Result<()> { // Validation reads an index's name, column, and field id, never its HNSW // tuning, so the writer's build params are not needed here. - let index_configs = build_index_configs(dataset, index_names, &HashMap::new()).await?; + let index_configs = build_index_configs( + dataset, + index_names, + &HashMap::new(), + OnMissingIndex::Reject, + ) + .await?; // The shard schema is base + `_tombstone`, as `ShardWriter::open` extends // it; field ids and the primary key resolve against that, not the base. From 568c525596cd6e4173179d391ac0753981b62584 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 15:28:02 -0400 Subject: [PATCH 07/48] fix(mem_wal): reconcile union arms across a schema change UnionExec requires its arms to agree on schema and does not reconcile. A sealed generation holds the shape it was sealed under, so once the base table's schema moves the LSM scan fails: an added or renamed column errors and a retyped one aborts the process. Bring every arm to the base table's shape first - cast where the type moved, fill typed nulls where the column is absent - and project each SSTable with its own generation's schema. --- .../src/dataset/mem_wal/scanner/planner.rs | 50 +++++++++++++++++-- .../src/dataset/mem_wal/scanner/projection.rs | 39 ++++++++++----- 2 files changed, 71 insertions(+), 18 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 2a83447fb17..75aae58a7b8 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -228,12 +228,40 @@ impl LsmScanPlanner { scan }; - source_plans.push(plan); + source_plans.push((plan, is_base)); } + // Every arm has to agree before the union: a generation is written under + // the schema the shard held when it was sealed, so one sealed before a + // column was added does not carry it, one sealed before a rename carries + // the old name, and one sealed before a retype carries the old type. + // `UnionExec` requires schema equality and does not reconcile. + // + // The base arm is the authority when it is here -- it is the only source + // the schema change was applied to -- and the newest generation + // otherwise. + let target = source_plans + .iter() + .find(|(_, is_base)| *is_base) + .or_else(|| source_plans.last()) + .map(|(plan, _)| plan.schema()); + let mut source_plans = match target { + Some(target) => source_plans + .into_iter() + .map(|(plan, _)| { + if plan.schema() == target { + Ok(plan) + } else { + project_to_canonical(plan, &target) + } + }) + .collect::>>()?, + None => Vec::new(), + }; + // Union, then coalesce into a single partition (UnionExec emits one // per arm; downstream consumers only read partition 0). - let mut plan: Arc = if source_plans.len() == 1 { + let plan: Arc = if source_plans.len() == 1 { source_plans.remove(0) } else { #[allow(deprecated)] @@ -243,7 +271,7 @@ impl LsmScanPlanner { // Project to the canonical output schema, dropping `_rowaddr` / // `_memtable_gen` unless the caller opted in. - plan = project_to_canonical( + let mut plan = project_to_canonical( plan, &self.canonical_scan_schema(projection, with_memtable_gen, keep_row_address), )?; @@ -332,9 +360,21 @@ impl LsmScanPlanner { .await?; let mut scanner = dataset.scan(); + // Asked of this generation, not of the base table. A generation + // is written under the schema the shard held when it was sealed, + // so a column added since is not in it and cannot be projected + // from it -- the arms are reconciled above the union instead. + // Projecting the base table's columns here asks an older file + // for a column it has never had, failing the scan outright. + let generation_schema: SchemaRef = Arc::new(dataset.schema().into()); let cols = - build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + build_scanner_projection(projection, &generation_schema, &self.pk_columns); + let cols: Vec<&str> = cols + .iter() + .filter(|c| generation_schema.column_with_name(c).is_some()) + .map(|s| s.as_str()) + .collect(); + scanner.project(&cols)?; scanner.with_row_address(); // Drop tombstones: fold `NOT _tombstone` into the predicate so diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index fd52ac90c0b..dcab0d9157b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs @@ -19,13 +19,14 @@ use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_expr::expressions::{Column, Literal}; +use datafusion::physical_expr::expressions::{CastExpr, Column, Literal}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::scalar::ScalarValue; use lance_core::{ROW_ADDR, ROW_ID, Result, is_system_column}; use super::exec::SchemaRelabelExec; +use crate::dataset::mem_wal::TOMBSTONE; /// Column name for distance in vector search results. pub const DISTANCE_COLUMN: &str = "_distance"; @@ -221,20 +222,32 @@ pub fn project_to_canonical( for field in target_schema.fields() { let name = field.name(); let expr: Arc = match input_schema.column_with_name(name) { - Some((idx, _)) => Arc::new(Column::new(name, idx)), + Some((idx, source)) if source.data_type() == field.data_type() => { + Arc::new(Column::new(name, idx)) + } + // A generation sealed before this column was retyped carries the + // type it was sealed under. + Some((idx, _)) => Arc::new(CastExpr::new( + Arc::new(Column::new(name, idx)), + field.data_type().clone(), + None, + )), None if is_system_column(name) => Arc::new(Literal::new(ScalarValue::UInt64(None))), None if name == DISTANCE_COLUMN => Arc::new(Literal::new(ScalarValue::Float32(None))), - None => { - return Err(lance_core::Error::internal(format!( - "Column '{}' missing from canonical projection source schema (have: {:?})", - name, - input_schema - .fields() - .iter() - .map(|f| f.name().clone()) - .collect::>() - ))); - } + // A generation that predates deletes carries no tombstone column; + // its rows are all live. The column is non-nullable, so a null here + // fails the scan outright. + None if name == TOMBSTONE => Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))), + // A generation sealed before this column existed, or under the name + // it carried then. Typed nulls are what it holds for those rows. + None => Arc::new(Literal::new( + ScalarValue::try_from(field.data_type()).map_err(|e| { + lance_core::Error::internal(format!( + "no null literal for column '{name}' of type {}: {e}", + field.data_type() + )) + })?, + )), }; project_exprs.push((expr, name.clone())); } From 9d3d03d49ab2a0e468a25565efa0d3217e53b436 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 16:38:20 -0400 Subject: [PATCH 08/48] fix(mem_wal): refuse a computed add_columns while a MemWAL is attached A computed transform derives the new column from the rows it can read, which is the committed fragments and nothing else. Rows still held in the WAL are invisible to it, so they take a null and keep it when they merge down - a wrong value written silently, that no later pass corrects. Emptiness is not a safe exemption: a write can land between the check and the commit, so the only race-free rule is to refuse whenever a MemWAL is present. AllNulls stays allowed, since a null is what it means everywhere. --- rust/lance/src/dataset/schema_evolution.rs | 32 ++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 8c1977d1605..f929a869c58 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -12,6 +12,7 @@ use super::{ transaction::{Operation, Transaction}, write::cleanup_data_fragments, }; +use crate::dataset::mem_wal::DatasetMemWalExt; use crate::index::load_all_indices; use crate::{Error, Result, io::exec::Planner}; use arrow::compute::CastOptions; @@ -446,6 +447,8 @@ pub(super) async fn add_columns( read_columns: Option>, batch_size: Option, ) -> Result<()> { + let computed = !matches!(transforms, NewColumnTransform::AllNulls(_)); + reject_computed_transform_on_mem_wal(dataset, computed).await?; let (fragments, schema, _fragments_to_cleanup, preserves_nullability) = add_columns_to_fragments( dataset, @@ -471,6 +474,35 @@ pub(super) async fn add_columns( .await } +/// Refuse a computed `add_columns` on a table with a MemWAL attached. +/// +/// A computed transform derives the new column from rows it can read, which is +/// the committed fragments and nothing else. Rows still in the WAL are invisible +/// to it, so they acquire the column as a null and keep that null when they +/// merge down -- a wrong value, written silently, that no later pass corrects. +/// +/// Emptiness is not a safe exemption: a write may land between the check and the +/// commit, so the only race-free rule is to refuse whenever a MemWAL is present. +/// `AllNulls` is exempt because a null is what it means everywhere. +/// +/// Takes the decision as a bool rather than the transform itself: a reference to +/// `NewColumnTransform` held across an await would require it to be `Sync`, +/// which its boxed reader is not. +async fn reject_computed_transform_on_mem_wal(dataset: &Dataset, computed: bool) -> Result<()> { + if !computed { + return Ok(()); + } + if dataset.mem_wal_index_details().await?.is_none() { + return Ok(()); + } + Err(Error::invalid_input( + "cannot add a computed column to a table with a MemWAL attached: rows held \ + in the WAL are not visible to the transform and would take a null. Add the \ + column as all-nulls and backfill it, or drop the MemWAL first." + .to_string(), + )) +} + async fn cleanup_new_column_data_files(fragments: &[FileFragment], new_fragments: &[Fragment]) { let Some(first_fragment) = fragments.first() else { return; From 3cdcd20016cb5cc21f6123a11a5486d3eca98f92 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 17:25:00 -0400 Subject: [PATCH 09/48] Revert "fix(mem_wal): refuse a computed add_columns while a MemWAL is attached" Refusing removes the feature rather than fixing it: a WAL-backed table could no longer take a computed or UDF column at all. The fix belongs where the WAL can actually be drained - sophon seals and compacts the fresh tier into base before materialising, so the transform sees every committed row. This reverts commit aeb8496ad4d2ee4b38c85e9e8df1a1c8a1e29df1. --- rust/lance/src/dataset/schema_evolution.rs | 32 ---------------------- 1 file changed, 32 deletions(-) diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index f929a869c58..8c1977d1605 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -12,7 +12,6 @@ use super::{ transaction::{Operation, Transaction}, write::cleanup_data_fragments, }; -use crate::dataset::mem_wal::DatasetMemWalExt; use crate::index::load_all_indices; use crate::{Error, Result, io::exec::Planner}; use arrow::compute::CastOptions; @@ -447,8 +446,6 @@ pub(super) async fn add_columns( read_columns: Option>, batch_size: Option, ) -> Result<()> { - let computed = !matches!(transforms, NewColumnTransform::AllNulls(_)); - reject_computed_transform_on_mem_wal(dataset, computed).await?; let (fragments, schema, _fragments_to_cleanup, preserves_nullability) = add_columns_to_fragments( dataset, @@ -474,35 +471,6 @@ pub(super) async fn add_columns( .await } -/// Refuse a computed `add_columns` on a table with a MemWAL attached. -/// -/// A computed transform derives the new column from rows it can read, which is -/// the committed fragments and nothing else. Rows still in the WAL are invisible -/// to it, so they acquire the column as a null and keep that null when they -/// merge down -- a wrong value, written silently, that no later pass corrects. -/// -/// Emptiness is not a safe exemption: a write may land between the check and the -/// commit, so the only race-free rule is to refuse whenever a MemWAL is present. -/// `AllNulls` is exempt because a null is what it means everywhere. -/// -/// Takes the decision as a bool rather than the transform itself: a reference to -/// `NewColumnTransform` held across an await would require it to be `Sync`, -/// which its boxed reader is not. -async fn reject_computed_transform_on_mem_wal(dataset: &Dataset, computed: bool) -> Result<()> { - if !computed { - return Ok(()); - } - if dataset.mem_wal_index_details().await?.is_none() { - return Ok(()); - } - Err(Error::invalid_input( - "cannot add a computed column to a table with a MemWAL attached: rows held \ - in the WAL are not visible to the transform and would take a null. Add the \ - column as all-nulls and backfill it, or drop the MemWAL first." - .to_string(), - )) -} - async fn cleanup_new_column_data_files(fragments: &[FileFragment], new_fragments: &[Fragment]) { let Some(first_fragment) = fragments.first() else { return; From 35732ae185a447ba35e69de5d5508405c12dba99 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 17:38:17 -0400 Subject: [PATCH 10/48] docs(mem_wal): say what the code is, in this repo's vocabulary Name the moment the maintained set is fixed as `initialize_mem_wal`, drop an inline comment that restated the doc above it, and state the id-before-name rule without arguing the alternatives. --- rust/lance/src/dataset/mem_wal/api.rs | 12 ++++++------ rust/lance/src/dataset/mem_wal/write.rs | 17 +++++------------ 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index ae39bd30a07..9387acfbf32 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -750,14 +750,14 @@ async fn build_index_configs( .into_iter() .next(); - // An index the maintained set names and the dataset no longer has: + // An index the maintained set names and the dataset does not have: // dropped outright, or carried away with the column it covered. The set - // is fixed when the write spec is installed and cannot be edited - // afterwards, so refusing here refuses the claim -- and a table whose - // claim cannot be built serves no reads at all, for a condition that - // costs only the fresh tier's copy of one index. + // is fixed at `initialize_mem_wal` and cannot be edited afterwards, so + // refusing here refuses the claim -- and a table whose claim cannot be + // built serves no reads at all, for a condition that costs only the + // fresh tier's copy of one index. // - // Serve without it instead. The base index is gone for everyone; the + // Serve without it. The base index is gone for everyone; the // fresh tier simply has nothing to keep in step with. let Some(index_meta) = index_meta else { if on_missing == OnMissingIndex::Reject { diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 949eea193f5..8bbddc5c86a 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -1625,17 +1625,14 @@ fn field_id_of(field: &ArrowField) -> Option { /// typed nulls; `_tombstone` is filled with `false`. A column the batch carries /// and the schema does not declare is dropped. /// -/// Ids are tried first because a rename keeps the id and changes the name, so a -/// name match would null the new name and drop the old one — the column's values -/// lost. Entries written before ids were carried have none, and fall back to the -/// name match, which is why both tiers exist. +/// Ids are tried first because a rename keeps the id and changes the name: a +/// name match would null the new name and drop the old one, losing the column's +/// values. An entry carrying no ids falls back to the name match. /// /// Both cases are what a replayed WAL entry looks like after the table's schema /// moved: an entry predates a column added since, and carries one dropped -/// since. Matching by position instead would reject the first outright and -/// store the second under its neighbour's name. Live writes reach here already -/// checked against the logical schema, so for them every column is present and -/// this only appends `_tombstone`. +/// since. Live writes reach here already checked against the logical schema, so +/// for them every column is present and this only appends `_tombstone`. /// /// A column whose type changed is cast, the same way `alter_columns` casts the /// base data, so a replayed row lands in the state it would have had if it had @@ -1650,10 +1647,6 @@ fn conform_to_storage_schema( pk_columns: &[String], ) -> Result { let n = batch.num_rows(); - // A field id survives a rename; a name does not. Matching on it first is - // what keeps a renamed column's values attached to the column, instead of - // nulling the new name and dropping the old one. Absent on entries written - // before ids were carried, which is why the name match below remains. let by_field_id: HashMap = batch .schema() .fields() From eecf43301c45d4eb26f245a77b982199791a7e41 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 19:32:46 -0400 Subject: [PATCH 11/48] fix(mem_wal): keep a renamed column's values across the union The union matches arms by name. A generation sealed under the old name stores the column under that name, so the match found nothing and the column was filled with typed nulls -- every row still present, every value gone. Relabel each generation's columns to the names the base table uses for the same field ids before the union sees them. A rename keeps the id, so the values arrive under the name the other arms use. Columns the base table does not declare keep their own names, as do all of them when no base arm is there to agree with. --- .../src/dataset/mem_wal/scanner/planner.rs | 83 ++++++++++++++++++- 1 file changed, 80 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 75aae58a7b8..b1da4550517 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -3,6 +3,7 @@ //! Query planner for LSM scanner. +use std::collections::HashMap; use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef}; @@ -11,13 +12,17 @@ use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; use lance_core::Result; +use lance_core::datatypes::Schema as LanceSchema; use tracing::instrument; use crate::dataset::mem_wal::TOMBSTONE; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; -use super::exec::{MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN}; +use super::exec::{ + MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, + SchemaRelabelExec, +}; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, validate_projection_names, @@ -55,6 +60,70 @@ pub struct LsmScanPlanner { warmer: Option>, } +/// What the base table calls each of its field ids, if a base arm is present. +/// +/// A rename changes a field's name and keeps its id, so this is what turns a +/// generation's stored names into the names every other arm uses. +fn base_field_names(sources: &[LsmDataSource]) -> Option> { + sources.iter().find_map(|source| match source { + LsmDataSource::BaseTable { dataset } => Some( + dataset + .schema() + .fields + .iter() + .map(|f| (f.id, f.name.clone())) + .collect(), + ), + _ => None, + }) +} + +/// Rename `plan`'s output columns to the names the base table uses for the same +/// field ids. +/// +/// Matched by id, so a column renamed since this generation was sealed keeps +/// its values instead of arriving under a name no other arm has. Columns the +/// base table does not declare -- `_rowaddr`, `_tombstone` -- keep their own +/// names, as do all of them when there is no base arm to agree with. +fn relabel_to_base_names( + plan: Arc, + generation_schema: &LanceSchema, + base_names: Option<&HashMap>, +) -> Arc { + let Some(base_names) = base_names else { + return plan; + }; + let ids: HashMap<&str, i32> = generation_schema + .fields + .iter() + .map(|f| (f.name.as_str(), f.id)) + .collect(); + let schema = plan.schema(); + let mut renamed = false; + let fields: Vec = schema + .fields() + .iter() + .map(|field| { + let base_name = ids + .get(field.name().as_str()) + .and_then(|id| base_names.get(id)) + .filter(|name| *name != field.name()); + match base_name { + Some(name) => { + renamed = true; + field.as_ref().clone().with_name(name) + } + None => field.as_ref().clone(), + } + }) + .collect(); + if !renamed { + return plan; + } + let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); + Arc::new(SchemaRelabelExec::new(plan, schema)) +} + impl LsmScanPlanner { /// Create a new planner. pub fn new( @@ -173,6 +242,12 @@ impl LsmScanPlanner { // down safely. The active memtable is in-memory and is never capped. let n_needed = limit.map(|l| l.saturating_add(offset.unwrap_or(0))); + // What the base table calls each field id today. A generation sealed + // under an older name still stores the column under that name, and the + // union below matches arms by name, so the generation's columns are + // relabelled to these before they meet the base arm. + let base_names = base_field_names(&sources); + let mut source_plans = Vec::new(); for source in sources { let is_base = matches!(source, LsmDataSource::BaseTable { .. }); @@ -186,7 +261,7 @@ impl LsmScanPlanner { _ => None, }; let scan = self - .build_source_scan(&source, projection, filter, fetch) + .build_source_scan(&source, projection, filter, fetch, base_names.as_ref()) .await?; // Drop cross-generation stale rows (PKs superseded by a newer gen). @@ -323,6 +398,7 @@ impl LsmScanPlanner { projection: Option<&[String]>, filter: Option<&Expr>, fetch: Option, + base_names: Option<&HashMap>, ) -> Result> { match source { LsmDataSource::BaseTable { dataset } => { @@ -399,7 +475,8 @@ impl LsmScanPlanner { scanner.limit(Some(fetch as i64), None)?; } - scanner.create_plan().await + let plan = scanner.create_plan().await?; + Ok(relabel_to_base_names(plan, dataset.schema(), base_names)) } LsmDataSource::ActiveMemTable { batch_store, From bb6242d43ecbb795813d9bfc1dbca741baa35941 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 22:49:38 -0400 Subject: [PATCH 12/48] fix(mem_wal): keep a renamed nested field's values across the union A struct's child can be renamed on its own: the parent column keeps its name, so the arms agree there, and only the struct's inner name differs. Casting one struct to the other is refused for having no field-name overlap, which failed the whole scan. Match nested fields by id too, and rebuild the struct array when a child is renamed -- the child names live in the array's own type, not in the schema above it, so relabelling the schema alone leaves the batch disagreeing with it. The children are reused, so the rebuild costs a pointer copy. --- .../mem_wal/scanner/exec/schema_relabel.rs | 46 +++++++++++- .../src/dataset/mem_wal/scanner/planner.rs | 72 +++++++++++-------- 2 files changed, 87 insertions(+), 31 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index 88691f99d63..88246ad2f2b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -8,8 +8,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use arrow_array::{RecordBatch, RecordBatchOptions}; -use arrow_schema::SchemaRef; +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, StructArray}; +use arrow_schema::{DataType, SchemaRef}; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::EquivalenceProperties; @@ -123,9 +123,18 @@ impl Stream for SchemaRelabelStream { Poll::Ready(Some(Ok(batch))) => { // Carry the row count explicitly: `try_new` infers it from the // first column, which a column-less batch does not have. + // A struct's child names live in the array's own type, not in + // the schema above it, so a renamed child needs the array + // rebuilt. The children are reused, so it costs a pointer copy. + let columns: Vec = batch + .columns() + .iter() + .zip(self.schema.fields()) + .map(|(column, field)| relabel_array(column, field.data_type())) + .collect(); let relabeled = RecordBatch::try_new_with_options( self.schema.clone(), - batch.columns().to_vec(), + columns, &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), ) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)); @@ -243,3 +252,34 @@ mod tests { ); } } + +/// `array` with its type relabelled to `target`, rebuilding a struct whose +/// children are named differently. +/// +/// Only names differ here -- the layout is identical -- so an array whose type +/// already matches, and any type this cannot express, is returned as it is and +/// left for `RecordBatch::try_new` to reject. +fn relabel_array(array: &ArrayRef, target: &DataType) -> ArrayRef { + if array.data_type() == target { + return Arc::clone(array); + } + let (DataType::Struct(_), DataType::Struct(target_fields)) = (array.data_type(), target) else { + return Arc::clone(array); + }; + let Some(source) = array.as_any().downcast_ref::() else { + return Arc::clone(array); + }; + if source.columns().len() != target_fields.len() { + return Arc::clone(array); + } + let children: Vec = source + .columns() + .iter() + .zip(target_fields) + .map(|(child, field)| relabel_array(child, field.data_type())) + .collect(); + match StructArray::try_new(target_fields.clone(), children, source.nulls().cloned()) { + Ok(rebuilt) => Arc::new(rebuilt), + Err(_) => Arc::clone(array), + } +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index b1da4550517..b15c9778a6f 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -12,7 +12,7 @@ use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; use lance_core::Result; -use lance_core::datatypes::Schema as LanceSchema; +use lance_core::datatypes::{Field as LanceField, Schema as LanceSchema}; use tracing::instrument; use crate::dataset::mem_wal::TOMBSTONE; @@ -63,21 +63,54 @@ pub struct LsmScanPlanner { /// What the base table calls each of its field ids, if a base arm is present. /// /// A rename changes a field's name and keeps its id, so this is what turns a -/// generation's stored names into the names every other arm uses. +/// generation's stored names into the names every other arm uses. Nested fields +/// are included: a struct's child can be renamed on its own, which changes the +/// parent's Arrow type without changing the parent's name. fn base_field_names(sources: &[LsmDataSource]) -> Option> { + fn collect(fields: &[LanceField], into: &mut HashMap) { + for field in fields { + into.insert(field.id, field.name.clone()); + collect(&field.children, into); + } + } sources.iter().find_map(|source| match source { - LsmDataSource::BaseTable { dataset } => Some( - dataset - .schema() - .fields - .iter() - .map(|f| (f.id, f.name.clone())) - .collect(), - ), + LsmDataSource::BaseTable { dataset } => { + let mut names = HashMap::new(); + collect(&dataset.schema().fields, &mut names); + Some(names) + } _ => None, }) } +/// One field renamed to the base table's name for its id, recursing into a +/// struct's children so a renamed child is matched by its own id. +fn rename_field( + field: &Field, + generation_fields: &[LanceField], + base_names: &HashMap, + renamed: &mut bool, +) -> Field { + let Some(source) = generation_fields.iter().find(|f| f.name == *field.name()) else { + return field.as_ref().clone(); + }; + let field = match base_names.get(&source.id).filter(|n| *n != field.name()) { + Some(name) => { + *renamed = true; + field.as_ref().clone().with_name(name) + } + None => field.as_ref().clone(), + }; + let DataType::Struct(children) = field.data_type() else { + return field; + }; + let children: Vec = children + .iter() + .map(|child| rename_field(child, &source.children, base_names, renamed)) + .collect(); + field.with_data_type(DataType::Struct(children.into())) +} + /// Rename `plan`'s output columns to the names the base table uses for the same /// field ids. /// @@ -93,29 +126,12 @@ fn relabel_to_base_names( let Some(base_names) = base_names else { return plan; }; - let ids: HashMap<&str, i32> = generation_schema - .fields - .iter() - .map(|f| (f.name.as_str(), f.id)) - .collect(); let schema = plan.schema(); let mut renamed = false; let fields: Vec = schema .fields() .iter() - .map(|field| { - let base_name = ids - .get(field.name().as_str()) - .and_then(|id| base_names.get(id)) - .filter(|name| *name != field.name()); - match base_name { - Some(name) => { - renamed = true; - field.as_ref().clone().with_name(name) - } - None => field.as_ref().clone(), - } - }) + .map(|field| rename_field(field, &generation_schema.fields, base_names, &mut renamed)) .collect(); if !renamed { return plan; From 51b508906dad24477b0a18048d994e276b8ac1c6 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 23:26:04 -0400 Subject: [PATCH 13/48] fix(mem_wal): never relabel a generation-only column onto a base name A column the base table does not declare exists only in a generation, so its id is drawn from the generation's own numbering and can collide with the id base gave an unrelated column -- relabelling by that id moves `_tombstone` onto a user column's name. Skip a rename onto a name the arm already carries. --- .../src/dataset/mem_wal/scanner/planner.rs | 33 +++++++++++++++---- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index b15c9778a6f..69415f96277 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -3,7 +3,7 @@ //! Query planner for LSM scanner. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef}; @@ -89,12 +89,17 @@ fn rename_field( field: &Field, generation_fields: &[LanceField], base_names: &HashMap, + taken: &HashSet<&str>, renamed: &mut bool, ) -> Field { let Some(source) = generation_fields.iter().find(|f| f.name == *field.name()) else { return field.as_ref().clone(); }; - let field = match base_names.get(&source.id).filter(|n| *n != field.name()) { + let rename = base_names + .get(&source.id) + .filter(|n| *n != field.name()) + .filter(|n| !taken.contains(n.as_str())); + let field = match rename { Some(name) => { *renamed = true; field.as_ref().clone().with_name(name) @@ -106,7 +111,7 @@ fn rename_field( }; let children: Vec = children .iter() - .map(|child| rename_field(child, &source.children, base_names, renamed)) + .map(|child| rename_field(child, &source.children, base_names, taken, renamed)) .collect(); field.with_data_type(DataType::Struct(children.into())) } @@ -115,9 +120,14 @@ fn rename_field( /// field ids. /// /// Matched by id, so a column renamed since this generation was sealed keeps -/// its values instead of arriving under a name no other arm has. Columns the -/// base table does not declare -- `_rowaddr`, `_tombstone` -- keep their own -/// names, as do all of them when there is no base arm to agree with. +/// its values instead of arriving under a name no other arm has. All of them +/// keep their own names when there is no base arm to agree with. +/// +/// A column the base table does not declare -- `_tombstone`, `_rowaddr` -- is +/// never renamed onto a base name: it exists only in a generation, so its id is +/// drawn from the generation's own numbering and can collide with the id base +/// gave an unrelated column. A rename onto a name the arm already carries is +/// skipped for the same reason. fn relabel_to_base_names( plan: Arc, generation_schema: &LanceSchema, @@ -127,11 +137,20 @@ fn relabel_to_base_names( return plan; }; let schema = plan.schema(); + let taken: HashSet<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); let mut renamed = false; let fields: Vec = schema .fields() .iter() - .map(|field| rename_field(field, &generation_schema.fields, base_names, &mut renamed)) + .map(|field| { + rename_field( + field, + &generation_schema.fields, + base_names, + &taken, + &mut renamed, + ) + }) .collect(); if !renamed { return plan; From 287d3caef46b08f051537040f928b5d7a9684cea Mon Sep 17 00:00:00 2001 From: XYZhan Date: Mon, 14 Sep 2026 23:40:59 -0400 Subject: [PATCH 14/48] fix(mem_wal): never rename a generation-only column onto a base name A generation's `_tombstone` is numbered in the generation's own schema, so its id collides with whatever base gave that id -- relabelling by id carried the tombstone in under a user column's name, where a column added after the seal read as that tombstone cast to its type instead of as null. Skipping a rename onto a name the arm already carries does not cover it: the name the tombstone was taking is one the arm does not have. --- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 69415f96277..ae313642713 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -13,6 +13,7 @@ use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; use lance_core::Result; use lance_core::datatypes::{Field as LanceField, Schema as LanceSchema}; +use lance_core::is_system_column; use tracing::instrument; use crate::dataset::mem_wal::TOMBSTONE; @@ -95,8 +96,9 @@ fn rename_field( let Some(source) = generation_fields.iter().find(|f| f.name == *field.name()) else { return field.as_ref().clone(); }; - let rename = base_names - .get(&source.id) + let rename = (field.name() != TOMBSTONE && !is_system_column(field.name())) + .then(|| base_names.get(&source.id)) + .flatten() .filter(|n| *n != field.name()) .filter(|n| !taken.contains(n.as_str())); let field = match rename { @@ -124,10 +126,10 @@ fn rename_field( /// keep their own names when there is no base arm to agree with. /// /// A column the base table does not declare -- `_tombstone`, `_rowaddr` -- is -/// never renamed onto a base name: it exists only in a generation, so its id is -/// drawn from the generation's own numbering and can collide with the id base -/// gave an unrelated column. A rename onto a name the arm already carries is -/// skipped for the same reason. +/// never renamed: it exists only in a generation, so its id is drawn from the +/// generation's own numbering and collides with whatever base gave that id, +/// which would carry a tombstone in under a user column's name. A rename onto +/// a name the arm already carries is skipped for the same reason. fn relabel_to_base_names( plan: Arc, generation_schema: &LanceSchema, From 51c6a5e4b972cda053142b6bd0bb8781e2ded5e9 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 00:05:38 -0400 Subject: [PATCH 15/48] test(mem_wal): cover the relabel's collision rules directly Each rule earned its place from a defect: a column takes base's name for its id, a column whose id base no longer declares keeps its own, a generation-only column is never renamed onto a base name, a rename onto a name the arm carries is skipped, and a struct child is matched by its own id. --- .../src/dataset/mem_wal/scanner/planner.rs | 120 ++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index ae313642713..dfabb34d9a1 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -573,6 +573,126 @@ mod tests { use super::*; use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; + /// A lance field, as a generation's schema records it. + fn lance_field(name: &str, id: i32, children: Vec) -> LanceField { + let arrow = if children.is_empty() { + Field::new(name, DataType::Int64, true) + } else { + Field::new( + name, + DataType::Struct( + children + .iter() + .map(|c| Field::new(&c.name, DataType::Int64, true)) + .collect(), + ), + true, + ) + }; + let mut field = LanceField::try_from(&arrow).expect("lance field"); + field.set_id(-1, &mut (id - 1).clone()); + field.id = id; + for (child, source) in field.children.iter_mut().zip(children.iter()) { + child.id = source.id; + } + field + } + + fn names(pairs: &[(i32, &str)]) -> HashMap { + pairs.iter().map(|(id, n)| (*id, n.to_string())).collect() + } + + fn rename_one( + field: &Field, + generation: &[LanceField], + base: &HashMap, + taken: &[&str], + ) -> (Field, bool) { + let taken: HashSet<&str> = taken.iter().copied().collect(); + let mut renamed = false; + let out = rename_field(field, generation, base, &taken, &mut renamed); + (out, renamed) + } + + #[test] + fn a_column_takes_the_name_base_now_gives_its_id() { + let generation = vec![lance_field("value", 1, vec![])]; + let base = names(&[(1, "amount")]); + let field = Field::new("value", DataType::Int64, true); + let (out, renamed) = rename_one(&field, &generation, &base, &["id", "value"]); + assert_eq!(out.name(), "amount"); + assert!(renamed); + } + + #[test] + fn a_column_base_no_longer_declares_keeps_its_name() { + // A retype gives the column a new id, so its old id is absent from base. + let generation = vec![lance_field("value", 1, vec![])]; + let base = names(&[(2, "value")]); + let field = Field::new("value", DataType::Int64, true); + let (out, renamed) = rename_one(&field, &generation, &base, &["value"]); + assert_eq!(out.name(), "value"); + assert!(!renamed); + } + + #[test] + fn a_tombstone_is_never_renamed_onto_a_base_name() { + // `_tombstone` is numbered in the generation's own schema, so its id + // collides with whatever base gave that id. Renaming it would carry a + // tombstone in under a user column's name. + let generation = vec![lance_field(TOMBSTONE, 2, vec![])]; + let base = names(&[(2, "extra")]); + let field = Field::new(TOMBSTONE, DataType::Boolean, false); + let (out, renamed) = rename_one(&field, &generation, &base, &[TOMBSTONE]); + assert_eq!(out.name(), TOMBSTONE); + assert!(!renamed, "a tombstone must keep its own name"); + } + + #[test] + fn a_rename_onto_a_name_the_arm_already_carries_is_skipped() { + let generation = vec![ + lance_field("value", 1, vec![]), + lance_field("other", 2, vec![]), + ]; + let base = names(&[(1, "other")]); + let field = Field::new("value", DataType::Int64, true); + let (out, renamed) = rename_one(&field, &generation, &base, &["value", "other"]); + assert_eq!(out.name(), "value", "renaming would collide with `other`"); + assert!(!renamed); + } + + #[test] + fn a_renamed_struct_child_is_matched_by_its_own_id() { + let generation = vec![lance_field("info", 1, vec![lance_field("c", 2, vec![])])]; + let base = names(&[(1, "info"), (2, "d")]); + let field = Field::new( + "info", + DataType::Struct(vec![Field::new("c", DataType::Int64, true)].into()), + true, + ); + let (out, renamed) = rename_one(&field, &generation, &base, &["id", "info"]); + assert!(renamed); + let DataType::Struct(children) = out.data_type() else { + panic!("expected a struct"); + }; + assert_eq!( + children[0].name(), + "d", + "the child takes base's name for its id" + ); + assert_eq!(out.name(), "info", "the parent's name has not moved"); + } + + #[test] + fn a_column_the_generation_does_not_declare_is_left_alone() { + let generation = vec![lance_field("value", 1, vec![])]; + let base = names(&[(1, "amount")]); + let field = Field::new("_rowaddr", DataType::UInt64, true); + let (out, renamed) = rename_one(&field, &generation, &base, &["_rowaddr"]); + assert_eq!(out.name(), "_rowaddr"); + assert!(!renamed); + } + fn create_test_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), From afa7e0d45521dcd16035fa1a75cfe7f8040cc329 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 12:02:08 -0400 Subject: [PATCH 16/48] fix(mem_wal): resolve a generation's columns by identity before reading them A generation stores each column under the name the table had when it was sealed, and every name the caller supplies is base's. Matching one to the other after the fact left four ways to answer wrongly: a named projection asked the file for a name it does not have and read nulls; a predicate did the same and failed the scan; a column renamed into a name a dropped column had returned the dropped column's data; and replay copied a renamed column into a new one that reused its name. Resolve the mapping first. By field id, which a rename keeps and a drop retires; then by name for a retype, which keeps the name and takes a new id, and only where no id match has claimed that name. A column base has dropped is not read at all. The caller's names are then translated into the generation's before projecting, and a predicate runs above the relabel -- with the columns it names filled in -- whenever this generation cannot evaluate it as written. The limit follows the filter rather than preceding it. --- .../src/dataset/mem_wal/scanner/planner.rs | 439 ++++++++++++------ rust/lance/src/dataset/mem_wal/write.rs | 9 +- 2 files changed, 308 insertions(+), 140 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index dfabb34d9a1..165cb8f4292 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -7,10 +7,14 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion::common::DFSchema; +use datafusion::execution::context::ExecutionProps; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; +use datafusion_physical_expr::create_physical_expr; use lance_core::Result; use lance_core::datatypes::{Field as LanceField, Schema as LanceSchema}; use lance_core::is_system_column; @@ -84,81 +88,204 @@ fn base_field_names(sources: &[LsmDataSource]) -> Option> { }) } -/// One field renamed to the base table's name for its id, recursing into a -/// struct's children so a renamed child is matched by its own id. -fn rename_field( - field: &Field, - generation_fields: &[LanceField], - base_names: &HashMap, - taken: &HashSet<&str>, - renamed: &mut bool, -) -> Field { - let Some(source) = generation_fields.iter().find(|f| f.name == *field.name()) else { - return field.as_ref().clone(); - }; - let rename = (field.name() != TOMBSTONE && !is_system_column(field.name())) - .then(|| base_names.get(&source.id)) - .flatten() - .filter(|n| *n != field.name()) - .filter(|n| !taken.contains(n.as_str())); - let field = match rename { - Some(name) => { - *renamed = true; - field.as_ref().clone().with_name(name) +/// How one generation's columns line up with the base table's. +/// +/// A generation stores each column under the name the table had when it was +/// sealed. Resolving that against base by field id -- which a rename keeps and +/// a drop retires -- gives three groups: columns base still declares, under the +/// name it uses now; columns base has dropped, which no longer belong in a read; +/// and columns that are the generation's own, like `_tombstone`, which base +/// never declared and which keep their names. +struct GenerationMapping { + /// Generation name → the name base uses for the same field id. + to_base: HashMap, + /// This generation's fields, for renaming a struct's children by their own + /// ids: a child can be renamed while its parent's name does not move. + fields: Vec, + /// What base calls each field id, children included. + base_names: HashMap, +} + +impl GenerationMapping { + fn resolve(generation_schema: &LanceSchema, base_names: Option<&HashMap>) -> Self { + let mut to_base = HashMap::new(); + let Some(base_names) = base_names else { + // No base arm to agree with, so this generation's own names are the + // names: every column maps to itself and none is retired. + for field in &generation_schema.fields { + to_base.insert(field.name.clone(), field.name.clone()); + } + return Self { + to_base, + fields: generation_schema.fields.clone(), + base_names: HashMap::new(), + }; + }; + // The generation's own columns are not base's to name or to retire. + let mine = |f: &LanceField| f.name == TOMBSTONE || is_system_column(&f.name); + + // By id first, which a rename keeps: these mappings are certain, and + // they are what makes the pass below able to tell the two remaining + // cases apart. + let mut claimed: HashSet<&str> = HashSet::new(); + for field in generation_schema.fields.iter().filter(|f| !mine(f)) { + if let Some(name) = base_names.get(&field.id) { + to_base.insert(field.name.clone(), name.clone()); + claimed.insert(name.as_str()); + } + } + + // What is left has an id base no longer declares, which happens two + // ways. A retype gives a column a new id while its name stays, so base + // still declares the name and no other column has claimed it -- the + // same column, matched by name. Otherwise the column is one base has + // dropped, and reading it would answer with data the table no longer + // has. + let base_has = |name: &str| base_names.values().any(|n| n == name); + for field in generation_schema.fields.iter().filter(|f| !mine(f)) { + if to_base.contains_key(&field.name) { + continue; + } + if base_has(&field.name) && !claimed.contains(field.name.as_str()) { + to_base.insert(field.name.clone(), field.name.clone()); + } + } + Self { + to_base, + fields: generation_schema.fields.clone(), + base_names: base_names.clone(), + } + } + + /// This generation's names for the base-table columns in `wanted`. + /// + /// A column the generation never had is simply absent from the result; the + /// union fills it in. A retired column is never returned even when its name + /// matches something base declares today. + fn generation_names_for(&self, wanted: &[String]) -> Vec { + let from_base: HashMap<&str, &str> = self + .to_base + .iter() + .map(|(stored, base)| (base.as_str(), stored.as_str())) + .collect(); + wanted + .iter() + .filter_map(|name| from_base.get(name.as_str()).map(|s| s.to_string())) + .collect() + } + + /// Whether this generation can evaluate `expr` as written. + /// + /// True only when every column the predicate names is stored here under + /// that same name. Otherwise the predicate runs above the relabel, where + /// the columns have base's names and the ones this generation never had are + /// filled in. + fn can_evaluate(&self, expr: &Expr, generation_schema: &LanceSchema) -> bool { + expr.column_refs().iter().all(|column| { + generation_schema.field(&column.name).is_some() + && self + .to_base + .get(&column.name) + .is_none_or(|base| base == &column.name) + }) + } + + /// Rename `plan`'s output columns to the names base uses, a struct's + /// children included. + fn relabel(&self, plan: Arc) -> Arc { + let schema = plan.schema(); + let mut renamed = false; + let fields: Vec = schema + .fields() + .iter() + .map(|field| self.rename(field, &self.fields, &mut renamed)) + .collect(); + if !renamed { + return plan; } - None => field.as_ref().clone(), - }; - let DataType::Struct(children) = field.data_type() else { - return field; - }; - let children: Vec = children - .iter() - .map(|child| rename_field(child, &source.children, base_names, taken, renamed)) - .collect(); - field.with_data_type(DataType::Struct(children.into())) + let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); + Arc::new(SchemaRelabelExec::new(plan, schema)) + } + + /// One field under base's name for its id, recursing into a struct so a + /// renamed child is matched by its own id rather than by the parent's. + fn rename(&self, field: &Field, among: &[LanceField], renamed: &mut bool) -> Field { + let Some(source) = among.iter().find(|f| f.name == *field.name()) else { + return field.as_ref().clone(); + }; + let field = match self.base_names.get(&source.id) { + Some(base) if base != field.name() => { + *renamed = true; + field.as_ref().clone().with_name(base) + } + _ => field.as_ref().clone(), + }; + let DataType::Struct(children) = field.data_type() else { + return field; + }; + let children: Vec = children + .iter() + .map(|child| self.rename(child, &source.children, renamed)) + .collect(); + field.with_data_type(DataType::Struct(children.into())) + } + + /// Whether base has dropped `name`, so this generation must not read it. + /// + /// A column base still declares is mapped; one it has dropped is not, and + /// neither is a column that was never base's to begin with. + #[cfg(test)] + fn retires(&self, name: &str) -> bool { + !self.to_base.contains_key(name) + && name != TOMBSTONE + && !is_system_column(name) + && self.fields.iter().any(|f| f.name == name) + } } -/// Rename `plan`'s output columns to the names the base table uses for the same -/// field ids. -/// -/// Matched by id, so a column renamed since this generation was sealed keeps -/// its values instead of arriving under a name no other arm has. All of them -/// keep their own names when there is no base arm to agree with. +/// Add the columns `expr` names that `plan` does not carry, as typed nulls from +/// `base_schema`. /// -/// A column the base table does not declare -- `_tombstone`, `_rowaddr` -- is -/// never renamed: it exists only in a generation, so its id is drawn from the -/// generation's own numbering and collides with whatever base gave that id, -/// which would carry a tombstone in under a user column's name. A rename onto -/// a name the arm already carries is skipped for the same reason. -fn relabel_to_base_names( +/// A generation sealed before a column existed holds no value for it, and null +/// is what that means — so a predicate over it can be answered rather than +/// refused. +fn fill_missing( plan: Arc, - generation_schema: &LanceSchema, - base_names: Option<&HashMap>, -) -> Arc { - let Some(base_names) = base_names else { - return plan; - }; + expr: &Expr, + base_schema: &SchemaRef, +) -> Result> { let schema = plan.schema(); - let taken: HashSet<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - let mut renamed = false; - let fields: Vec = schema - .fields() - .iter() - .map(|field| { - rename_field( - field, - &generation_schema.fields, - base_names, - &taken, - &mut renamed, - ) - }) - .collect(); - if !renamed { - return plan; + let mut fields: Vec = schema.fields().iter().map(|f| f.as_ref().clone()).collect(); + let mut added = false; + for column in expr.column_refs() { + if schema.column_with_name(&column.name).is_some() { + continue; + } + let Ok(field) = base_schema.field_with_name(&column.name) else { + continue; + }; + fields.push(field.clone().with_nullable(true)); + added = true; + } + if !added { + return Ok(plan); } - let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); - Arc::new(SchemaRelabelExec::new(plan, schema)) + let target = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); + project_to_canonical(plan, &target) +} + +/// Apply `expr` above a plan whose columns have just been relabelled, for a +/// generation whose names the predicate could not be run against directly. +fn filter_above(plan: Arc, expr: &Expr) -> Result> { + let schema = plan.schema(); + let df_schema = DFSchema::try_from(schema.as_ref().clone()) + .map_err(|e| lance_core::Error::internal(format!("filter schema: {e}")))?; + let props = ExecutionProps::new(); + let physical = create_physical_expr(expr, &df_schema, &props) + .map_err(|e| lance_core::Error::internal(format!("plan filter `{expr}`: {e}")))?; + Ok(Arc::new(FilterExec::try_new(physical, plan).map_err( + |e| lance_core::Error::internal(format!("filter: {e}")), + )?)) } impl LsmScanPlanner { @@ -473,21 +600,22 @@ impl LsmScanPlanner { .await?; let mut scanner = dataset.scan(); - // Asked of this generation, not of the base table. A generation - // is written under the schema the shard held when it was sealed, - // so a column added since is not in it and cannot be projected - // from it -- the arms are reconciled above the union instead. - // Projecting the base table's columns here asks an older file - // for a column it has never had, failing the scan outright. - let generation_schema: SchemaRef = Arc::new(dataset.schema().into()); - let cols = - build_scanner_projection(projection, &generation_schema, &self.pk_columns); - let cols: Vec<&str> = cols - .iter() - .filter(|c| generation_schema.column_with_name(c).is_some()) - .map(|s| s.as_str()) - .collect(); - scanner.project(&cols)?; + // Which of this generation's columns the base table still + // declares, and what it calls each of them. Resolved by field + // id before anything else looks at a name: a rename changes the + // name and keeps the id, and a column whose id base no longer + // declares is a column the table has dropped. + let mapping = GenerationMapping::resolve(dataset.schema(), base_names); + + // Projected under this generation's own names, so an older file + // is asked only for columns it has. The caller names columns as + // the base table does, so those are translated first; one the + // generation never had is left out and filled in above the + // union. + let wanted = + build_scanner_projection(projection, &self.base_schema, &self.pk_columns); + let cols = mapping.generation_names_for(&wanted); + scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; scanner.with_row_address(); // Drop tombstones: fold `NOT _tombstone` into the predicate so @@ -495,25 +623,50 @@ impl LsmScanPlanner { // The older real row a tombstone supersedes is dropped by the // cross-gen block-list, not by this filter. Gen written before // deletes existed lack the column → no fold, nothing to drop. + // + // A caller's predicate names columns as base does, so it can + // only run here when this generation agrees with base on every + // name. Otherwise it runs above the relabel, and the limit goes + // with it -- a limit under an unapplied filter would cut rows + // the filter has not seen. + // The predicate names columns as base does. This generation + // can only evaluate it when it stores every one of them under + // that same name: a rename moved the name, and a column added + // since the seal is not here at all, though `IS NULL` over it + // is a question about these rows that still has an answer. + let evaluable = + filter.is_none_or(|expr| mapping.can_evaluate(expr, dataset.schema())); + let caller_filter = if evaluable { filter } else { None }; let folded; let effective: Option<&Expr> = if dataset.schema().field(TOMBSTONE).is_some() { - folded = fold_not_tombstone(filter); + folded = fold_not_tombstone(caller_filter); Some(&folded) } else { - filter + caller_filter }; if let Some(expr) = effective { scanner.filter_expr(expr.clone()); } - // Per-source limit pushdown: SSTables are - // within-gen live (dedup-on-flush deletion vectors), so any - // `fetch` post-filter rows are valid contributions. + // A limit under a filter that has not run would cut rows the + // filter never saw. if let Some(fetch) = fetch { - scanner.limit(Some(fetch as i64), None)?; + if evaluable { + scanner.limit(Some(fetch as i64), None)?; + } } let plan = scanner.create_plan().await?; - Ok(relabel_to_base_names(plan, dataset.schema(), base_names)) + let plan = mapping.relabel(plan); + match filter { + Some(expr) if !evaluable => { + // The predicate may name a column this generation never + // had, and `IS NULL` over it is a question about these + // rows with a real answer. Fill those in before asking. + let plan = fill_missing(plan, expr, &self.base_schema)?; + filter_above(plan, expr) + } + _ => Ok(plan), + } } LsmDataSource::ActiveMemTable { batch_store, @@ -602,75 +755,83 @@ mod tests { pairs.iter().map(|(id, n)| (*id, n.to_string())).collect() } - fn rename_one( - field: &Field, - generation: &[LanceField], - base: &HashMap, - taken: &[&str], - ) -> (Field, bool) { - let taken: HashSet<&str> = taken.iter().copied().collect(); - let mut renamed = false; - let out = rename_field(field, generation, base, &taken, &mut renamed); - (out, renamed) + fn mapping(generation: Vec, base: &[(i32, &str)]) -> GenerationMapping { + let schema = LanceSchema { + fields: generation, + metadata: Default::default(), + }; + GenerationMapping::resolve(&schema, Some(&names(base))) } #[test] fn a_column_takes_the_name_base_now_gives_its_id() { - let generation = vec![lance_field("value", 1, vec![])]; - let base = names(&[(1, "amount")]); - let field = Field::new("value", DataType::Int64, true); - let (out, renamed) = rename_one(&field, &generation, &base, &["id", "value"]); - assert_eq!(out.name(), "amount"); - assert!(renamed); + let m = mapping(vec![lance_field("value", 1, vec![])], &[(1, "amount")]); + assert_eq!(m.to_base.get("value").map(String::as_str), Some("amount")); + assert_eq!(m.generation_names_for(&["amount".into()]), vec!["value"]); + assert!(m.renames_projected(&["value".into()])); } #[test] - fn a_column_base_no_longer_declares_keeps_its_name() { - // A retype gives the column a new id, so its old id is absent from base. - let generation = vec![lance_field("value", 1, vec![])]; - let base = names(&[(2, "value")]); - let field = Field::new("value", DataType::Int64, true); - let (out, renamed) = rename_one(&field, &generation, &base, &["value"]); - assert_eq!(out.name(), "value"); - assert!(!renamed); + fn a_column_base_no_longer_declares_is_retired() { + // Dropping a column retires its id, so a generation still holding it + // must not contribute it -- not even to a base column that has since + // taken the name. + let m = mapping( + vec![ + lance_field("value", 1, vec![]), + lance_field("other", 2, vec![]), + ], + &[(1, "other")], + ); + assert!(m.retires("other"), "id 2 is gone from base"); + assert_eq!( + m.generation_names_for(&["other".into()]), + vec!["value"], + "base's `other` is this generation's `value`, by id" + ); } #[test] - fn a_tombstone_is_never_renamed_onto_a_base_name() { + fn a_generation_only_column_is_neither_renamed_nor_retired() { // `_tombstone` is numbered in the generation's own schema, so its id - // collides with whatever base gave that id. Renaming it would carry a - // tombstone in under a user column's name. - let generation = vec![lance_field(TOMBSTONE, 2, vec![])]; - let base = names(&[(2, "extra")]); - let field = Field::new(TOMBSTONE, DataType::Boolean, false); - let (out, renamed) = rename_one(&field, &generation, &base, &[TOMBSTONE]); - assert_eq!(out.name(), TOMBSTONE); - assert!(!renamed, "a tombstone must keep its own name"); + // collides with whatever base gave that id. + let m = mapping( + vec![ + lance_field("value", 1, vec![]), + lance_field(TOMBSTONE, 2, vec![]), + ], + &[(1, "value"), (2, "extra")], + ); + assert!(!m.retires(TOMBSTONE)); + assert_eq!(m.to_base.get(TOMBSTONE), None); } #[test] - fn a_rename_onto_a_name_the_arm_already_carries_is_skipped() { - let generation = vec![ - lance_field("value", 1, vec![]), - lance_field("other", 2, vec![]), - ]; - let base = names(&[(1, "other")]); - let field = Field::new("value", DataType::Int64, true); - let (out, renamed) = rename_one(&field, &generation, &base, &["value", "other"]); - assert_eq!(out.name(), "value", "renaming would collide with `other`"); - assert!(!renamed); + fn a_column_the_generation_never_had_is_not_projected() { + let m = mapping( + vec![lance_field("value", 1, vec![])], + &[(1, "value"), (2, "added")], + ); + assert_eq!( + m.generation_names_for(&["value".into(), "added".into()]), + vec!["value"], + "the column added since is filled in above the union" + ); } #[test] fn a_renamed_struct_child_is_matched_by_its_own_id() { - let generation = vec![lance_field("info", 1, vec![lance_field("c", 2, vec![])])]; - let base = names(&[(1, "info"), (2, "d")]); + let m = mapping( + vec![lance_field("info", 1, vec![lance_field("c", 2, vec![])])], + &[(1, "info"), (2, "d")], + ); let field = Field::new( "info", DataType::Struct(vec![Field::new("c", DataType::Int64, true)].into()), true, ); - let (out, renamed) = rename_one(&field, &generation, &base, &["id", "info"]); + let mut renamed = false; + let out = m.rename(&field, &m.fields.clone(), &mut renamed); assert!(renamed); let DataType::Struct(children) = out.data_type() else { panic!("expected a struct"); @@ -685,10 +846,10 @@ mod tests { #[test] fn a_column_the_generation_does_not_declare_is_left_alone() { - let generation = vec![lance_field("value", 1, vec![])]; - let base = names(&[(1, "amount")]); + let m = mapping(vec![lance_field("value", 1, vec![])], &[(1, "amount")]); let field = Field::new("_rowaddr", DataType::UInt64, true); - let (out, renamed) = rename_one(&field, &generation, &base, &["_rowaddr"]); + let mut renamed = false; + let out = m.rename(&field, &m.fields.clone(), &mut renamed); assert_eq!(out.name(), "_rowaddr"); assert!(!renamed); } diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 8bbddc5c86a..0047235daf3 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -1660,7 +1660,14 @@ fn conform_to_storage_schema( let name = field.name(); let carried = field_id_of(field) .and_then(|id| by_field_id.get(&id).map(|c| (*c).clone())) - .or_else(|| batch.column_by_name(name).cloned()); + .or_else(|| match batch.schema().column_with_name(name) { + // The entry carries this name under a different id, so it is a + // different column: a rename freed the name and something else + // took it. Absent, not carried. + Some((_, f)) if field_id_of(f).is_some() && field_id_of(field).is_some() => None, + Some((i, _)) => Some(batch.column(i).clone()), + None => None, + }); if let Some(column) = carried { columns.push(if column.data_type() == field.data_type() { column From 85f6c175589f029063ed80d680c061a82a4b6f47 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 13:06:47 -0400 Subject: [PATCH 17/48] fix(mem_wal): let one stored column answer to one schema column A rename frees a name, so a stored column can answer to a schema column's name without being that column, and a cast takes a new field id, so the same column can answer to no id at all. Matching by id and then by name resolved both, but nothing stopped a single stored column from being taken twice -- once by an id match and again by a name match it was no longer entitled to. A name match now applies only to a column no id match has claimed. Replay and the scan resolve it the same way, so a retype keeps its values and a column that merely reuses a freed name reads as null. Three more, all in the scan: a predicate this generation cannot evaluate runs above the relabel, so the columns it names have to be read even when the caller did not ask for them; a column base has dropped is no longer a column a predicate may be answered from; and the fallback target is the newest generation, which is the first of them, since sources arrive generation-DESC. --- .../mem_wal/scanner/exec/schema_relabel.rs | 62 +++++++++---------- .../src/dataset/mem_wal/scanner/planner.rs | 47 +++++++++----- rust/lance/src/dataset/mem_wal/write.rs | 32 ++++++---- 3 files changed, 83 insertions(+), 58 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index 88246ad2f2b..1e525a24efc 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -151,6 +151,37 @@ impl datafusion::physical_plan::RecordBatchStream for SchemaRelabelStream { } } +/// `array` with its type relabelled to `target`, rebuilding a struct whose +/// children are named differently. +/// +/// Only names differ here -- the layout is identical -- so an array whose type +/// already matches, and any type this cannot express, is returned as it is and +/// left for `RecordBatch::try_new` to reject. +fn relabel_array(array: &ArrayRef, target: &DataType) -> ArrayRef { + if array.data_type() == target { + return Arc::clone(array); + } + let (DataType::Struct(_), DataType::Struct(target_fields)) = (array.data_type(), target) else { + return Arc::clone(array); + }; + let Some(source) = array.as_any().downcast_ref::() else { + return Arc::clone(array); + }; + if source.columns().len() != target_fields.len() { + return Arc::clone(array); + } + let children: Vec = source + .columns() + .iter() + .zip(target_fields) + .map(|(child, field)| relabel_array(child, field.data_type())) + .collect(); + match StructArray::try_new(target_fields.clone(), children, source.nulls().cloned()) { + Ok(rebuilt) => Arc::new(rebuilt), + Err(_) => Arc::clone(array), + } +} + #[cfg(test)] mod tests { use super::*; @@ -252,34 +283,3 @@ mod tests { ); } } - -/// `array` with its type relabelled to `target`, rebuilding a struct whose -/// children are named differently. -/// -/// Only names differ here -- the layout is identical -- so an array whose type -/// already matches, and any type this cannot express, is returned as it is and -/// left for `RecordBatch::try_new` to reject. -fn relabel_array(array: &ArrayRef, target: &DataType) -> ArrayRef { - if array.data_type() == target { - return Arc::clone(array); - } - let (DataType::Struct(_), DataType::Struct(target_fields)) = (array.data_type(), target) else { - return Arc::clone(array); - }; - let Some(source) = array.as_any().downcast_ref::() else { - return Arc::clone(array); - }; - if source.columns().len() != target_fields.len() { - return Arc::clone(array); - } - let children: Vec = source - .columns() - .iter() - .zip(target_fields) - .map(|(child, field)| relabel_array(child, field.data_type())) - .collect(); - match StructArray::try_new(target_fields.clone(), children, source.nulls().cloned()) { - Ok(rebuilt) => Arc::new(rebuilt), - Err(_) => Arc::clone(array), - } -} diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 165cb8f4292..65506cac716 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -182,11 +182,13 @@ impl GenerationMapping { /// filled in. fn can_evaluate(&self, expr: &Expr, generation_schema: &LanceSchema) -> bool { expr.column_refs().iter().all(|column| { - generation_schema.field(&column.name).is_some() - && self - .to_base - .get(&column.name) - .is_none_or(|base| base == &column.name) + // A positive mapping, not merely a column of that name: one base + // has dropped is still stored here, and answering from it would + // filter on data the table no longer has. + self.to_base + .get(&column.name) + .is_some_and(|base| base == &column.name) + && generation_schema.field(&column.name).is_some() }) } @@ -477,12 +479,12 @@ impl LsmScanPlanner { // `UnionExec` requires schema equality and does not reconcile. // // The base arm is the authority when it is here -- it is the only source - // the schema change was applied to -- and the newest generation - // otherwise. + // the schema change was applied to. Otherwise the newest generation is, + // and sources arrive generation-DESC, so it is the first of them. let target = source_plans .iter() .find(|(_, is_base)| *is_base) - .or_else(|| source_plans.last()) + .or_else(|| source_plans.first()) .map(|(plan, _)| plan.schema()); let mut source_plans = match target { Some(target) => source_plans @@ -612,8 +614,26 @@ impl LsmScanPlanner { // the base table does, so those are translated first; one the // generation never had is left out and filled in above the // union. - let wanted = + // The predicate names columns as base does. This generation can + // only evaluate it when it stores every one of them under that + // same name: a rename moved the name, and a column added since + // the seal is not here at all, though `IS NULL` over it is a + // question about these rows that still has an answer. + let evaluable = + filter.is_none_or(|expr| mapping.can_evaluate(expr, dataset.schema())); + + let mut wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); + // A filter this generation cannot evaluate runs above the + // relabel, reading the columns from this scan -- so they have to + // be in it, whether or not the caller asked for them. + if let Some(expr) = filter.filter(|_| !evaluable) { + for column in expr.column_refs() { + if !wanted.contains(&column.name) { + wanted.push(column.name.clone()); + } + } + } let cols = mapping.generation_names_for(&wanted); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; scanner.with_row_address(); @@ -634,8 +654,6 @@ impl LsmScanPlanner { // that same name: a rename moved the name, and a column added // since the seal is not here at all, though `IS NULL` over it // is a question about these rows that still has an answer. - let evaluable = - filter.is_none_or(|expr| mapping.can_evaluate(expr, dataset.schema())); let caller_filter = if evaluable { filter } else { None }; let folded; let effective: Option<&Expr> = if dataset.schema().field(TOMBSTONE).is_some() { @@ -649,10 +667,8 @@ impl LsmScanPlanner { } // A limit under a filter that has not run would cut rows the // filter never saw. - if let Some(fetch) = fetch { - if evaluable { - scanner.limit(Some(fetch as i64), None)?; - } + if let Some(fetch) = fetch.filter(|_| evaluable) { + scanner.limit(Some(fetch as i64), None)?; } let plan = scanner.create_plan().await?; @@ -768,7 +784,6 @@ mod tests { let m = mapping(vec![lance_field("value", 1, vec![])], &[(1, "amount")]); assert_eq!(m.to_base.get("value").map(String::as_str), Some("amount")); assert_eq!(m.generation_names_for(&["amount".into()]), vec!["value"]); - assert!(m.renames_projected(&["value".into()])); } #[test] diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 0047235daf3..2e3a2094b11 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -12,7 +12,7 @@ //! - [`IndexStore`] - In-memory index management //! - [`MemTableFlusher`] - Flush MemTable to storage as single Lance file -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt::Debug; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, RwLock as StdRwLock}; @@ -1647,26 +1647,36 @@ fn conform_to_storage_schema( pk_columns: &[String], ) -> Result { let n = batch.num_rows(); - let by_field_id: HashMap = batch + let by_field_id: HashMap = batch .schema() .fields() .iter() .enumerate() - .filter_map(|(i, f)| field_id_of(f).map(|id| (id, batch.column(i)))) + .filter_map(|(i, f)| field_id_of(f).map(|id| (id, i))) + .collect(); + + // An entry column an id match has already taken. A rename frees a name for + // something else to use, so the entry's column can answer to the schema's + // name without being the schema's column -- and it is the id match that + // says which one it really is. + let claimed: HashSet = storage_schema + .fields() + .iter() + .filter_map(|f| field_id_of(f)) + .filter_map(|id| by_field_id.get(&id).copied()) .collect(); let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); for field in storage_schema.fields() { let name = field.name(); let carried = field_id_of(field) - .and_then(|id| by_field_id.get(&id).map(|c| (*c).clone())) - .or_else(|| match batch.schema().column_with_name(name) { - // The entry carries this name under a different id, so it is a - // different column: a rename freed the name and something else - // took it. Absent, not carried. - Some((_, f)) if field_id_of(f).is_some() && field_id_of(field).is_some() => None, - Some((i, _)) => Some(batch.column(i).clone()), - None => None, + .and_then(|id| by_field_id.get(&id).map(|i| batch.column(*i).clone())) + .or_else(|| { + // No column of this id. The entry may still carry this column + // under this name with an id of its own -- a cast takes a new + // id and keeps the name -- but only if nothing has claimed it. + let (i, _) = batch.schema().column_with_name(name)?; + (!claimed.contains(&i)).then(|| batch.column(i).clone()) }); if let Some(column) = carried { columns.push(if column.data_type() == field.data_type() { From 64fabb15bc1bb1a544fa798348403a39ef8ddb4d Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 15:16:53 -0400 Subject: [PATCH 18/48] fix(mem_wal): resolve a schema change through one plan, and refuse a retype A cast takes a new field id and keeps the column's name, which is exactly what dropping a column and adding another under that name looks like. Nothing in the schemas tells the two apart, so every attempt to reconcile both ended up guessing: matching by name resurrected dropped values, matching by id alone lost retyped ones. A table with a MemWAL now refuses to change a column's type, before any of the request commits and including a request that also renames. A table without one is unaffected. With that gone, identity is the field id and nothing else, and the guessing goes with it. Replay and the scan resolve a schema change the same way, through one plan that says for each target column which stored column supplies it, which children need rebuilding, and what a column the source never had produces. Field ids are carried recursively, so a struct's children are matched by their own identity rather than by an Arrow cast that follows names. The target is the table's schema, supplied by the caller, rather than whichever source the collector ordered first. Projection and predicate are planned together: the columns a deferred predicate names are read even when the caller did not ask for them, and a limit follows a filter that has not run yet. --- rust/lance/src/dataset/mem_wal.rs | 48 +- rust/lance/src/dataset/mem_wal/reconcile.rs | 286 ++++++++++ .../src/dataset/mem_wal/scanner/builder.rs | 17 +- .../lance/src/dataset/mem_wal/scanner/exec.rs | 2 + .../dataset/mem_wal/scanner/exec/reconcile.rs | 128 +++++ .../src/dataset/mem_wal/scanner/planner.rs | 530 ++++++------------ rust/lance/src/dataset/mem_wal/write.rs | 111 +--- rust/lance/src/dataset/schema_evolution.rs | 26 + 8 files changed, 661 insertions(+), 487 deletions(-) create mode 100644 rust/lance/src/dataset/mem_wal/reconcile.rs create mode 100644 rust/lance/src/dataset/mem_wal/scanner/exec/reconcile.rs diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index cf4f856a2b0..ed376b7a5bf 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -38,6 +38,7 @@ pub mod index; mod manifest; pub mod memtable; pub mod observer; +pub(crate) mod reconcile; pub mod scanner; pub mod sharding; #[cfg(test)] @@ -48,9 +49,7 @@ pub mod write; use std::sync::Arc; -use std::collections::HashMap; - -use lance_core::datatypes::{LANCE_FIELD_ID_KEY, Schema}; +use lance_core::datatypes::{Field, LANCE_FIELD_ID_KEY, Schema}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; @@ -138,28 +137,39 @@ pub fn relax_non_pk_nullability( /// callers that compare schemas for equality. pub(crate) fn arrow_schema_with_field_ids(schema: &Schema) -> ArrowSchema { let arrow: ArrowSchema = schema.into(); - let ids: HashMap<&str, i32> = schema - .fields - .iter() - .map(|f| (f.name.as_str(), f.id)) - .collect(); let fields: Vec = arrow .fields() .iter() - .map(|field| { - let Some(id) = ids - .get(field.name().as_str()) - .copied() - .filter(|id| *id >= 0) - else { - return field.as_ref().clone(); - }; + .map(|field| stamp_field_id(field, &schema.fields)) + .collect(); + ArrowSchema::new_with_metadata(fields, arrow.metadata().clone()) +} + +/// One field carrying its lance id, and its struct children carrying theirs. +/// +/// A struct's children are fields in their own right: they have ids, a rename +/// moves one child's name and not the parent's, and a reader that cannot see a +/// child's id has only its name to go on. +fn stamp_field_id(field: &ArrowField, among: &[Field]) -> ArrowField { + let Some(source) = among.iter().find(|f| f.name == *field.name()) else { + return field.clone(); + }; + let field = match source.id { + id if id >= 0 => { let mut metadata = field.metadata().clone(); metadata.insert(LANCE_FIELD_ID_KEY.to_string(), id.to_string()); - field.as_ref().clone().with_metadata(metadata) - }) + field.clone().with_metadata(metadata) + } + _ => field.clone(), + }; + let DataType::Struct(children) = field.data_type() else { + return field; + }; + let children: Vec = children + .iter() + .map(|child| stamp_field_id(child, &source.children)) .collect(); - ArrowSchema::new_with_metadata(fields, arrow.metadata().clone()) + field.with_data_type(DataType::Struct(children.into())) } pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs new file mode 100644 index 00000000000..82dd2fd1bcd --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -0,0 +1,286 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Bringing a batch written under one schema to the schema in force now. +//! +//! A MemWAL holds rows written under whatever schema the table had at the time, +//! and they are read and replayed against the schema it has now. Resolving one +//! to the other is done once, here, and the result drives both: replay applies +//! it to a WAL entry, and a scan applies it to a generation's batches. +//! +//! Columns are matched by **field id**. A rename changes a field's name and +//! keeps its id, so a name is not identity; a cast keeps the name and takes a +//! new id, so a name is not identity there either. Where the two disagree the +//! plan refuses rather than guesses -- see [`Resolution::Ambiguous`]. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions, StructArray}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; +use lance_core::datatypes::LANCE_FIELD_ID_KEY; +use lance_core::{Error, Result}; + +use super::TOMBSTONE; + +/// The lance field id an Arrow field carries, if it carries one. +pub(crate) fn field_id_of(field: &ArrowField) -> Option { + field + .metadata() + .get(LANCE_FIELD_ID_KEY) + .and_then(|v| v.parse::().ok()) + .filter(|id| *id >= 0) +} + +/// `schema` without the field ids, for comparing against what a caller sends. +/// +/// Ids belong to the stored schema, where identity has to survive a rename. A +/// caller's batch carries none, and Arrow compares a struct's children by their +/// full field -- metadata included -- so a stamped schema would reject it. +pub(crate) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { + fn strip(field: &ArrowField) -> ArrowField { + let mut metadata = field.metadata().clone(); + metadata.remove(LANCE_FIELD_ID_KEY); + let field = field.clone().with_metadata(metadata); + match field.data_type() { + DataType::Struct(children) => { + let children: Vec = children.iter().map(|c| strip(c)).collect(); + field.with_data_type(DataType::Struct(children.into())) + } + _ => field, + } + } + let fields: Vec = schema.fields().iter().map(|f| strip(f)).collect(); + ArrowSchema::new_with_metadata(fields, schema.metadata().clone()) +} + +/// Where one target column's values come from. +#[derive(Debug, Clone)] +enum Source { + /// The source column at this index, as it stands. + Take(usize), + /// The source column at this index, whose struct children need their own + /// resolution. + Nested(usize, Vec, DataType), + /// The source does not have this column: rows written before it existed + /// hold no value for it. + Null(DataType), + /// `_tombstone`, which a generation written before deletes existed does not + /// carry. Its rows are all live. + Live, +} + +/// One resolution of a source schema against a target schema. +pub struct Plan { + target: SchemaRef, + sources: Vec, + /// Whether the source already is the target, so applying changes nothing. + identity: bool, +} + +impl Plan { + /// Resolve `source` against `target`, or say why it cannot be done. + /// + /// `pk_columns` may not be filled with nulls: a row with no primary key + /// cannot be placed, so an absent one is an error rather than a null. + pub(crate) fn resolve( + source: &ArrowSchema, + target: &SchemaRef, + pk_columns: &[String], + ) -> Result { + // An id match takes its source column; a name match may then only take + // one nothing has claimed. A rename frees a name for another column to + // use, and it is the id that says which column is really which. + let claimed = claimed_by_id(source, target); + let sources = target + .fields() + .iter() + .map(|field| resolve_field(field, source.fields(), &claimed, pk_columns)) + .collect::>>()?; + let identity = source.fields() == target.fields(); + Ok(Self { + target: Arc::clone(target), + sources, + identity, + }) + } + + /// The schema a batch has after [`Self::apply`]. + pub(crate) fn target(&self) -> &SchemaRef { + &self.target + } + + /// Whether the source schema already is the target, so applying this plan + /// would produce the batch it was given. + pub(crate) fn is_identity(&self) -> bool { + self.identity + } + + /// `batch` under the target schema. + pub(crate) fn apply(&self, batch: &RecordBatch) -> Result { + let rows = batch.num_rows(); + let columns = self + .sources + .iter() + .zip(self.target.fields()) + .map(|(source, field)| take_column(source, batch.columns(), rows, field.name())) + .collect::>>()?; + RecordBatch::try_new_with_options( + Arc::clone(&self.target), + columns, + &RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .map_err(|e| Error::invalid_input(format!("reconcile a batch to the schema: {e}"))) + } +} + +/// Source columns an id match has taken, which a name match may not take again. +fn claimed_by_id(source: &ArrowSchema, target: &SchemaRef) -> Vec { + let by_id: HashMap = source + .fields() + .iter() + .enumerate() + .filter_map(|(i, f)| field_id_of(f).map(|id| (id, i))) + .collect(); + let mut claimed = vec![false; source.fields().len()]; + for field in target.fields() { + if let Some(i) = field_id_of(field).and_then(|id| by_id.get(&id)) { + claimed[*i] = true; + } + } + claimed +} + +fn resolve_field( + field: &ArrowField, + source_fields: &arrow_schema::Fields, + claimed: &[bool], + pk_columns: &[String], +) -> Result { + let name = field.name(); + let by_id = field_id_of(field).and_then(|id| { + source_fields + .iter() + .position(|f| field_id_of(f) == Some(id)) + }); + // Identity is the field id where both sides carry one. A name is not: a + // rename moves the name and leaves the id, so a source column of the same + // name under a *different* id is a different column -- one dropped and + // another added under its name, whose values the table no longer has. + // + // A name match is right only where identity is absent: a batch a caller has + // just handed in carries no ids, and neither does a schema supplied by a + // caller who has none to give. + let by_name = || { + source_fields + .iter() + .position(|f| f.name() == name) + .filter(|i| !claimed[*i]) + }; + let index = match field_id_of(field) { + // The target names an identity: only that identity answers for it, + // unless the source has none to be matched on. + Some(_) => by_id.or_else(|| { + source_fields + .iter() + .all(|f| field_id_of(f).is_none()) + .then(by_name) + .flatten() + }), + None => by_name(), + }; + + let Some(index) = index else { + if name == TOMBSTONE { + return Ok(Source::Live); + } + if pk_columns.iter().any(|c| c == name) { + return Err(Error::invalid_input(format!( + "batch is missing primary key column `{name}` declared by the schema" + ))); + } + return Ok(Source::Null(field.data_type().clone())); + }; + + let source = &source_fields[index]; + if source.data_type() == field.data_type() { + return Ok(Source::Take(index)); + } + // The same column under a different scalar type. A table with a MemWAL + // refuses a cast, so this is a disagreement to surface rather than paper + // over; a struct differing only in its children is handled below. + if !matches!( + (source.data_type(), field.data_type()), + (DataType::Struct(_), DataType::Struct(_)) + ) { + return Err(Error::invalid_input(format!( + "column `{name}` is stored as {} and the schema declares {}; a column's type \ + cannot change on a table with a MemWAL", + source.data_type(), + field.data_type() + ))); + } + // A struct whose children were renamed keeps the parent's name and id, so + // only the children differ. Resolve them the same way rather than casting, + // which Arrow matches by name. + if let (DataType::Struct(source_children), DataType::Struct(target_children)) = + (source.data_type(), field.data_type()) + { + let claimed = claimed_children(source_children, target_children); + let children = target_children + .iter() + .map(|child| resolve_field(child, source_children, &claimed, &[])) + .collect::>>()?; + return Ok(Source::Nested(index, children, field.data_type().clone())); + } + unreachable!("a non-struct type mismatch is rejected above") +} + +fn claimed_children(source: &arrow_schema::Fields, target: &arrow_schema::Fields) -> Vec { + let by_id: HashMap = source + .iter() + .enumerate() + .filter_map(|(i, f)| field_id_of(f).map(|id| (id, i))) + .collect(); + let mut claimed = vec![false; source.len()]; + for field in target { + if let Some(i) = field_id_of(field).and_then(|id| by_id.get(&id)) { + claimed[*i] = true; + } + } + claimed +} + +fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) -> Result { + match source { + Source::Take(i) => Ok(Arc::clone(&columns[*i])), + Source::Nested(i, children, to) => { + let DataType::Struct(target_children) = to else { + unreachable!("Nested is only built for a struct target"); + }; + let struct_array = columns[*i] + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input(format!("column `{name}` is not a struct")))?; + let built = children + .iter() + .zip(target_children) + .map(|(child, field)| { + take_column(child, struct_array.columns(), rows, field.name()) + }) + .collect::>>()?; + Ok(Arc::new( + StructArray::try_new( + target_children.clone(), + built, + struct_array.nulls().cloned(), + ) + .map_err(|e| { + Error::invalid_input(format!("rebuild struct column `{name}`: {e}")) + })?, + )) + } + Source::Null(ty) => Ok(arrow_array::new_null_array(ty, rows)), + Source::Live => Ok(Arc::new(BooleanArray::from(vec![false; rows]))), + } +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs index dbd5e9e4a79..8bd5ddb94bf 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs @@ -197,6 +197,9 @@ pub struct LsmScanner { /// Derived from the base dataset when one is present, otherwise supplied /// explicitly by [`Self::without_base_table`]. schema: SchemaRef, + /// [`Self::schema`] with each field's id, which is what resolves a + /// generation's stored columns to the table's. + identity_schema: SchemaRef, shard_snapshots: Vec, /// In-memory memtables by shard (active + frozen-awaiting-flush), so /// the scanner path carries frozen-undrained generations too. @@ -259,6 +262,9 @@ impl LsmScanner { // path-bound store binding. let store_params = base_table.store_params().map(derived_store_params); Self { + identity_schema: Arc::new(crate::dataset::mem_wal::arrow_schema_with_field_ids( + base_table.schema(), + )), base: BaseSource::Table(base_table), schema: Arc::new(arrow_schema), shard_snapshots, @@ -302,6 +308,10 @@ impl LsmScanner { pk_columns: Vec, ) -> Self { Self { + // Whatever identity the caller supplied. Pass a schema carrying + // field ids to have a generation's columns resolved by them; a + // plain one is matched by name, as it was before ids existed. + identity_schema: schema.clone(), base: BaseSource::PathOnly(base_path.into()), schema, shard_snapshots, @@ -719,7 +729,12 @@ impl LsmScanner { }); } - let mut planner = LsmScanPlanner::new(collector, self.pk_columns.clone(), base_schema); + let mut planner = LsmScanPlanner::new( + collector, + self.pk_columns.clone(), + base_schema, + Arc::clone(&self.identity_schema), + ); if let Some(session) = &self.session { planner = planner.with_session(session.clone()); } diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/scanner/exec.rs index 2fa0a3f514f..a353b10c27e 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec.rs @@ -19,6 +19,7 @@ mod first_by_pk; mod generation_tag; mod pk; mod pk_block_filter; +mod reconcile; mod schema_relabel; pub use bloom_guard::{BloomFilterGuardExec, compute_pk_hash_from_scalars}; @@ -30,4 +31,5 @@ pub use pk::{ validate_pk_types, }; pub use pk_block_filter::PkBlockFilterExec; +pub use reconcile::ReconcileExec; pub use schema_relabel::SchemaRelabelExec; diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/reconcile.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/reconcile.rs new file mode 100644 index 00000000000..5f9e9847f93 --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/reconcile.rs @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Applying a [`Plan`] to a source's batches. + +use std::fmt; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt}; + +use crate::dataset::mem_wal::reconcile::Plan; + +/// Brings one source's batches to the schema the scan reads in. +/// +/// The same [`Plan`] replay applies to a WAL entry, so a generation and an +/// entry written under the same schema are reconciled the same way. +pub struct ReconcileExec { + input: Arc, + plan: Arc, + properties: Arc, +} + +impl ReconcileExec { + pub fn new(input: Arc, plan: Arc) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(plan.target())), + input.properties().output_partitioning().clone(), + input.properties().emission_type, + input.properties().boundedness, + )); + Self { + input, + plan, + properties, + } + } +} + +impl fmt::Debug for ReconcileExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ReconcileExec") + } +} + +impl DisplayAs for ReconcileExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "ReconcileExec") + } +} + +impl ExecutionPlan for ReconcileExec { + fn name(&self) -> &str { + "ReconcileExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(self.plan.target()) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DFResult> { + Ok(Arc::new(Self::new( + children + .into_iter() + .next() + .ok_or_else(|| DataFusionError::Internal("ReconcileExec needs one child".into()))?, + Arc::clone(&self.plan), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DFResult { + Ok(Box::pin(ReconcileStream { + input: self.input.execute(partition, context)?, + plan: Arc::clone(&self.plan), + })) + } +} + +struct ReconcileStream { + input: SendableRecordBatchStream, + plan: Arc, +} + +impl Stream for ReconcileStream { + type Item = DFResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.input.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => Poll::Ready(Some( + self.plan + .apply(&batch) + .map_err(|e| DataFusionError::External(Box::new(e))), + )), + other => other, + } + } +} + +impl RecordBatchStream for ReconcileStream { + fn schema(&self) -> SchemaRef { + Arc::clone(self.plan.target()) + } +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 65506cac716..8964c3366a0 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -3,10 +3,10 @@ //! Query planner for LSM scanner. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; -use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; use datafusion::common::DFSchema; use datafusion::execution::context::ExecutionProps; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; @@ -16,17 +16,16 @@ use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; use datafusion_physical_expr::create_physical_expr; use lance_core::Result; -use lance_core::datatypes::{Field as LanceField, Schema as LanceSchema}; use lance_core::is_system_column; use tracing::instrument; -use crate::dataset::mem_wal::TOMBSTONE; +use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; +use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{ - MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, - SchemaRelabelExec, + MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, ReconcileExec, }; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, @@ -55,6 +54,10 @@ pub struct LsmScanPlanner { pk_columns: Vec, /// Schema of the base table. base_schema: SchemaRef, + /// The same schema with each field's id, which is what resolves a + /// generation's columns to the table's. Supplied by the caller rather than + /// read off whichever source happens to be present. + identity_schema: SchemaRef, /// Session threaded into SSTable opens (shared caches). session: Option>, /// Store params for opening SSTables, reusing the base dataset's store. @@ -65,229 +68,103 @@ pub struct LsmScanPlanner { warmer: Option>, } -/// What the base table calls each of its field ids, if a base arm is present. -/// -/// A rename changes a field's name and keeps its id, so this is what turns a -/// generation's stored names into the names every other arm uses. Nested fields -/// are included: a struct's child can be renamed on its own, which changes the -/// parent's Arrow type without changing the parent's name. -fn base_field_names(sources: &[LsmDataSource]) -> Option> { - fn collect(fields: &[LanceField], into: &mut HashMap) { - for field in fields { - into.insert(field.id, field.name.clone()); - collect(&field.children, into); - } - } - sources.iter().find_map(|source| match source { - LsmDataSource::BaseTable { dataset } => { - let mut names = HashMap::new(); - collect(&dataset.schema().fields, &mut names); - Some(names) - } - _ => None, - }) +/// Apply `expr` above a source that has been reconciled, for a generation whose +/// stored names the predicate could not be run against directly. +fn filter_above(plan: Arc, expr: &Expr) -> Result> { + let schema = plan.schema(); + let df_schema = DFSchema::try_from(schema.as_ref().clone()) + .map_err(|e| lance_core::Error::internal(format!("filter schema: {e}")))?; + let props = ExecutionProps::new(); + let physical = create_physical_expr(expr, &df_schema, &props) + .map_err(|e| lance_core::Error::internal(format!("plan filter `{expr}`: {e}")))?; + Ok(Arc::new(FilterExec::try_new(physical, plan).map_err( + |e| lance_core::Error::internal(format!("filter: {e}")), + )?)) } -/// How one generation's columns line up with the base table's. +/// What the table calls each of this generation's stored columns, keyed by the +/// stored name. /// -/// A generation stores each column under the name the table had when it was -/// sealed. Resolving that against base by field id -- which a rename keeps and -/// a drop retires -- gives three groups: columns base still declares, under the -/// name it uses now; columns base has dropped, which no longer belong in a read; -/// and columns that are the generation's own, like `_tombstone`, which base -/// never declared and which keep their names. -struct GenerationMapping { - /// Generation name → the name base uses for the same field id. - to_base: HashMap, - /// This generation's fields, for renaming a struct's children by their own - /// ids: a child can be renamed while its parent's name does not move. - fields: Vec, - /// What base calls each field id, children included. - base_names: HashMap, -} - -impl GenerationMapping { - fn resolve(generation_schema: &LanceSchema, base_names: Option<&HashMap>) -> Self { - let mut to_base = HashMap::new(); - let Some(base_names) = base_names else { - // No base arm to agree with, so this generation's own names are the - // names: every column maps to itself and none is retired. - for field in &generation_schema.fields { - to_base.insert(field.name.clone(), field.name.clone()); - } - return Self { - to_base, - fields: generation_schema.fields.clone(), - base_names: HashMap::new(), - }; - }; - // The generation's own columns are not base's to name or to retire. - let mine = |f: &LanceField| f.name == TOMBSTONE || is_system_column(&f.name); - - // By id first, which a rename keeps: these mappings are certain, and - // they are what makes the pass below able to tell the two remaining - // cases apart. - let mut claimed: HashSet<&str> = HashSet::new(); - for field in generation_schema.fields.iter().filter(|f| !mine(f)) { - if let Some(name) = base_names.get(&field.id) { - to_base.insert(field.name.clone(), name.clone()); - claimed.insert(name.as_str()); - } - } - - // What is left has an id base no longer declares, which happens two - // ways. A retype gives a column a new id while its name stays, so base - // still declares the name and no other column has claimed it -- the - // same column, matched by name. Otherwise the column is one base has - // dropped, and reading it would answer with data the table no longer - // has. - let base_has = |name: &str| base_names.values().any(|n| n == name); - for field in generation_schema.fields.iter().filter(|f| !mine(f)) { - if to_base.contains_key(&field.name) { - continue; - } - if base_has(&field.name) && !claimed.contains(field.name.as_str()) { - to_base.insert(field.name.clone(), field.name.clone()); - } - } - Self { - to_base, - fields: generation_schema.fields.clone(), - base_names: base_names.clone(), - } - } - - /// This generation's names for the base-table columns in `wanted`. - /// - /// A column the generation never had is simply absent from the result; the - /// union fills it in. A retired column is never returned even when its name - /// matches something base declares today. - fn generation_names_for(&self, wanted: &[String]) -> Vec { - let from_base: HashMap<&str, &str> = self - .to_base - .iter() - .map(|(stored, base)| (base.as_str(), stored.as_str())) - .collect(); - wanted +/// Matched by field id, which a rename keeps. A stored column whose id the +/// table no longer declares is absent from the result: the table has dropped +/// it, and reading it would answer with data the table no longer has. +fn stored_names(stored: &Schema, table: &Schema) -> HashMap { + eprintln!( + "DBG stored={:?} table={:?}", + stored + .fields() .iter() - .filter_map(|name| from_base.get(name.as_str()).map(|s| s.to_string())) - .collect() - } - - /// Whether this generation can evaluate `expr` as written. - /// - /// True only when every column the predicate names is stored here under - /// that same name. Otherwise the predicate runs above the relabel, where - /// the columns have base's names and the ones this generation never had are - /// filled in. - fn can_evaluate(&self, expr: &Expr, generation_schema: &LanceSchema) -> bool { - expr.column_refs().iter().all(|column| { - // A positive mapping, not merely a column of that name: one base - // has dropped is still stored here, and answering from it would - // filter on data the table no longer has. - self.to_base - .get(&column.name) - .is_some_and(|base| base == &column.name) - && generation_schema.field(&column.name).is_some() - }) - } - - /// Rename `plan`'s output columns to the names base uses, a struct's - /// children included. - fn relabel(&self, plan: Arc) -> Arc { - let schema = plan.schema(); - let mut renamed = false; - let fields: Vec = schema + .map(|f| (f.name().clone(), field_id_of(f))) + .collect::>(), + table .fields() .iter() - .map(|field| self.rename(field, &self.fields, &mut renamed)) - .collect(); - if !renamed { - return plan; - } - let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); - Arc::new(SchemaRelabelExec::new(plan, schema)) - } - - /// One field under base's name for its id, recursing into a struct so a - /// renamed child is matched by its own id rather than by the parent's. - fn rename(&self, field: &Field, among: &[LanceField], renamed: &mut bool) -> Field { - let Some(source) = among.iter().find(|f| f.name == *field.name()) else { - return field.as_ref().clone(); - }; - let field = match self.base_names.get(&source.id) { - Some(base) if base != field.name() => { - *renamed = true; - field.as_ref().clone().with_name(base) - } - _ => field.as_ref().clone(), - }; - let DataType::Struct(children) = field.data_type() else { - return field; - }; - let children: Vec = children + .map(|f| (f.name().clone(), field_id_of(f))) + .collect::>() + ); + let by_id: HashMap = table + .fields() + .iter() + .filter_map(|f| field_id_of(f).map(|id| (id, f.name().as_str()))) + .collect(); + // A caller that supplied no ids has only names to be matched on, which is + // how this worked before ids were carried at all. + if by_id.is_empty() { + return stored + .fields() .iter() - .map(|child| self.rename(child, &source.children, renamed)) + .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) + .filter(|f| table.field_with_name(f.name()).is_ok()) + .map(|f| (f.name().clone(), f.name().clone())) .collect(); - field.with_data_type(DataType::Struct(children.into())) - } - - /// Whether base has dropped `name`, so this generation must not read it. - /// - /// A column base still declares is mapped; one it has dropped is not, and - /// neither is a column that was never base's to begin with. - #[cfg(test)] - fn retires(&self, name: &str) -> bool { - !self.to_base.contains_key(name) - && name != TOMBSTONE - && !is_system_column(name) - && self.fields.iter().any(|f| f.name == name) } + stored + .fields() + .iter() + // A generation's own columns are numbered in its own schema, so their + // ids collide with whatever the table gave those numbers. They are not + // the table's columns and are never resolved to one. + .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) + .filter_map(|f| { + let id = field_id_of(f)?; + by_id + .get(&id) + .map(|name| (f.name().clone(), name.to_string())) + }) + .collect() } -/// Add the columns `expr` names that `plan` does not carry, as typed nulls from -/// `base_schema`. +/// `schema` with each field carrying the id `stored` gives the same name. /// -/// A generation sealed before a column existed holds no value for it, and null -/// is what that means — so a predicate over it can be answered rather than -/// refused. -fn fill_missing( - plan: Arc, - expr: &Expr, - base_schema: &SchemaRef, -) -> Result> { - let schema = plan.schema(); - let mut fields: Vec = schema.fields().iter().map(|f| f.as_ref().clone()).collect(); - let mut added = false; - for column in expr.column_refs() { - if schema.column_with_name(&column.name).is_some() { - continue; - } - let Ok(field) = base_schema.field_with_name(&column.name) else { - continue; +/// A scan's output schema is built from the dataset and carries no ids, so they +/// are put back before identity is resolved against it. +fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { + fn restore(field: &Field, among: &Fields) -> Field { + let Some(source) = among.iter().find(|f| f.name() == field.name()) else { + return field.clone(); }; - fields.push(field.clone().with_nullable(true)); - added = true; - } - if !added { - return Ok(plan); + let mut metadata = field.metadata().clone(); + metadata.extend(source.metadata().clone()); + let field = field.clone().with_metadata(metadata); + // A struct's children carry their own ids, and a child can be renamed + // while its parent's name does not move. + match (field.data_type(), source.data_type()) { + (DataType::Struct(children), DataType::Struct(source_children)) => { + let children: Vec = children + .iter() + .map(|child| restore(child, source_children)) + .collect(); + field.with_data_type(DataType::Struct(children.into())) + } + _ => field, + } } - let target = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())); - project_to_canonical(plan, &target) -} - -/// Apply `expr` above a plan whose columns have just been relabelled, for a -/// generation whose names the predicate could not be run against directly. -fn filter_above(plan: Arc, expr: &Expr) -> Result> { - let schema = plan.schema(); - let df_schema = DFSchema::try_from(schema.as_ref().clone()) - .map_err(|e| lance_core::Error::internal(format!("filter schema: {e}")))?; - let props = ExecutionProps::new(); - let physical = create_physical_expr(expr, &df_schema, &props) - .map_err(|e| lance_core::Error::internal(format!("plan filter `{expr}`: {e}")))?; - Ok(Arc::new(FilterExec::try_new(physical, plan).map_err( - |e| lance_core::Error::internal(format!("filter: {e}")), - )?)) + let fields: Vec = schema + .fields() + .iter() + .map(|field| restore(field, stored.fields())) + .collect(); + Schema::new_with_metadata(fields, schema.metadata().clone()) } impl LsmScanPlanner { @@ -296,11 +173,13 @@ impl LsmScanPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, + identity_schema: SchemaRef, ) -> Self { Self { collector, pk_columns, base_schema, + identity_schema, session: None, store_params: None, sstable_cache: None, @@ -408,12 +287,6 @@ impl LsmScanPlanner { // down safely. The active memtable is in-memory and is never capped. let n_needed = limit.map(|l| l.saturating_add(offset.unwrap_or(0))); - // What the base table calls each field id today. A generation sealed - // under an older name still stores the column under that name, and the - // union below matches arms by name, so the generation's columns are - // relabelled to these before they meet the base arm. - let base_names = base_field_names(&sources); - let mut source_plans = Vec::new(); for source in sources { let is_base = matches!(source, LsmDataSource::BaseTable { .. }); @@ -427,7 +300,7 @@ impl LsmScanPlanner { _ => None, }; let scan = self - .build_source_scan(&source, projection, filter, fetch, base_names.as_ref()) + .build_source_scan(&source, projection, filter, fetch) .await?; // Drop cross-generation stale rows (PKs superseded by a newer gen). @@ -558,13 +431,39 @@ impl LsmScanPlanner { } /// Build scan plan for a single data source. + /// What one generation's scan should produce once reconciled: the columns + /// this query needs, named as the table names them, plus the ones the + /// generation carries of its own. + /// + /// The same for every generation, and taken from the table rather than from + /// whichever source the collector ordered first. + fn generation_target( + &self, + wanted: &[String], + source: &Schema, + names: &HashMap, + ) -> SchemaRef { + let mut fields: Vec = wanted + .iter() + .filter_map(|name| self.identity_schema.field_with_name(name).ok().cloned()) + .collect(); + // A generation's own columns are not the table's, so they pass through + // as the generation has them. + for field in source.fields() { + let is_the_tables = names.contains_key(field.name()); + if !is_the_tables && fields.iter().all(|f| f.name() != field.name()) { + fields.push(field.as_ref().clone()); + } + } + Arc::new(Schema::new(fields)) + } + async fn build_source_scan( &self, source: &LsmDataSource, projection: Option<&[String]>, filter: Option<&Expr>, fetch: Option, - base_names: Option<&HashMap>, ) -> Result> { match source { LsmDataSource::BaseTable { dataset } => { @@ -602,40 +501,41 @@ impl LsmScanPlanner { .await?; let mut scanner = dataset.scan(); - // Which of this generation's columns the base table still - // declares, and what it calls each of them. Resolved by field - // id before anything else looks at a name: a rename changes the - // name and keeps the id, and a column whose id base no longer - // declares is a column the table has dropped. - let mapping = GenerationMapping::resolve(dataset.schema(), base_names); - - // Projected under this generation's own names, so an older file - // is asked only for columns it has. The caller names columns as - // the base table does, so those are translated first; one the - // generation never had is left out and filled in above the - // union. - // The predicate names columns as base does. This generation can - // only evaluate it when it stores every one of them under that - // same name: a rename moved the name, and a column added since - // the seal is not here at all, though `IS NULL` over it is a - // question about these rows that still has an answer. - let evaluable = - filter.is_none_or(|expr| mapping.can_evaluate(expr, dataset.schema())); + // What the table calls each of this generation's columns, by + // field id: a rename changes the name and keeps the id. + let stored = arrow_schema_with_field_ids(dataset.schema()); + let names = stored_names(&stored, &self.identity_schema); + // Asked of this generation under its own names, so an older + // file is only asked for columns it has. A column it never had + // is filled in after the scan. let mut wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - // A filter this generation cannot evaluate runs above the - // relabel, reading the columns from this scan -- so they have to - // be in it, whether or not the caller asked for them. - if let Some(expr) = filter.filter(|_| !evaluable) { + // A predicate this generation cannot answer as written runs + // after reconciliation, reading its columns from this scan -- + // so they have to be in it whether the caller asked or not. + let answerable = filter.is_none_or(|expr| { + expr.column_refs() + .iter() + .all(|c| names.get(c.name.as_str()) == Some(&c.name)) + }); + if let Some(expr) = filter.filter(|_| !answerable) { for column in expr.column_refs() { if !wanted.contains(&column.name) { wanted.push(column.name.clone()); } } } - let cols = mapping.generation_names_for(&wanted); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + let cols: Vec<&str> = wanted + .iter() + .filter_map(|name| { + names + .iter() + .find(|(_, table_name)| *table_name == name) + .map(|(stored_name, _)| stored_name.as_str()) + }) + .collect(); + scanner.project(&cols)?; scanner.with_row_address(); // Drop tombstones: fold `NOT _tombstone` into the predicate so @@ -643,18 +543,7 @@ impl LsmScanPlanner { // The older real row a tombstone supersedes is dropped by the // cross-gen block-list, not by this filter. Gen written before // deletes existed lack the column → no fold, nothing to drop. - // - // A caller's predicate names columns as base does, so it can - // only run here when this generation agrees with base on every - // name. Otherwise it runs above the relabel, and the limit goes - // with it -- a limit under an unapplied filter would cut rows - // the filter has not seen. - // The predicate names columns as base does. This generation - // can only evaluate it when it stores every one of them under - // that same name: a rename moved the name, and a column added - // since the seal is not here at all, though `IS NULL` over it - // is a question about these rows that still has an answer. - let caller_filter = if evaluable { filter } else { None }; + let caller_filter = if answerable { filter } else { None }; let folded; let effective: Option<&Expr> = if dataset.schema().field(TOMBSTONE).is_some() { folded = fold_not_tombstone(caller_filter); @@ -667,21 +556,22 @@ impl LsmScanPlanner { } // A limit under a filter that has not run would cut rows the // filter never saw. - if let Some(fetch) = fetch.filter(|_| evaluable) { + if let Some(fetch) = fetch.filter(|_| answerable) { scanner.limit(Some(fetch as i64), None)?; } - let plan = scanner.create_plan().await?; - let plan = mapping.relabel(plan); + let scan = scanner.create_plan().await?; + // Planned against what the scan actually produces, which is the + // projection plus `_rowaddr`, and carrying the ids the dataset + // gives those columns. + let source = with_ids_from(&scan.schema(), &stored); + let target = self.generation_target(&wanted, &source, &names); + let plan = Plan::resolve(&source, &target, &self.pk_columns)?; + let reconciled: Arc = + Arc::new(ReconcileExec::new(scan, Arc::new(plan))); match filter { - Some(expr) if !evaluable => { - // The predicate may name a column this generation never - // had, and `IS NULL` over it is a question about these - // rows with a real answer. Fill those in before asking. - let plan = fill_missing(plan, expr, &self.base_schema)?; - filter_above(plan, expr) - } - _ => Ok(plan), + Some(expr) if !answerable => filter_above(reconciled, expr), + _ => Ok(reconciled), } } LsmDataSource::ActiveMemTable { @@ -771,104 +661,6 @@ mod tests { pairs.iter().map(|(id, n)| (*id, n.to_string())).collect() } - fn mapping(generation: Vec, base: &[(i32, &str)]) -> GenerationMapping { - let schema = LanceSchema { - fields: generation, - metadata: Default::default(), - }; - GenerationMapping::resolve(&schema, Some(&names(base))) - } - - #[test] - fn a_column_takes_the_name_base_now_gives_its_id() { - let m = mapping(vec![lance_field("value", 1, vec![])], &[(1, "amount")]); - assert_eq!(m.to_base.get("value").map(String::as_str), Some("amount")); - assert_eq!(m.generation_names_for(&["amount".into()]), vec!["value"]); - } - - #[test] - fn a_column_base_no_longer_declares_is_retired() { - // Dropping a column retires its id, so a generation still holding it - // must not contribute it -- not even to a base column that has since - // taken the name. - let m = mapping( - vec![ - lance_field("value", 1, vec![]), - lance_field("other", 2, vec![]), - ], - &[(1, "other")], - ); - assert!(m.retires("other"), "id 2 is gone from base"); - assert_eq!( - m.generation_names_for(&["other".into()]), - vec!["value"], - "base's `other` is this generation's `value`, by id" - ); - } - - #[test] - fn a_generation_only_column_is_neither_renamed_nor_retired() { - // `_tombstone` is numbered in the generation's own schema, so its id - // collides with whatever base gave that id. - let m = mapping( - vec![ - lance_field("value", 1, vec![]), - lance_field(TOMBSTONE, 2, vec![]), - ], - &[(1, "value"), (2, "extra")], - ); - assert!(!m.retires(TOMBSTONE)); - assert_eq!(m.to_base.get(TOMBSTONE), None); - } - - #[test] - fn a_column_the_generation_never_had_is_not_projected() { - let m = mapping( - vec![lance_field("value", 1, vec![])], - &[(1, "value"), (2, "added")], - ); - assert_eq!( - m.generation_names_for(&["value".into(), "added".into()]), - vec!["value"], - "the column added since is filled in above the union" - ); - } - - #[test] - fn a_renamed_struct_child_is_matched_by_its_own_id() { - let m = mapping( - vec![lance_field("info", 1, vec![lance_field("c", 2, vec![])])], - &[(1, "info"), (2, "d")], - ); - let field = Field::new( - "info", - DataType::Struct(vec![Field::new("c", DataType::Int64, true)].into()), - true, - ); - let mut renamed = false; - let out = m.rename(&field, &m.fields.clone(), &mut renamed); - assert!(renamed); - let DataType::Struct(children) = out.data_type() else { - panic!("expected a struct"); - }; - assert_eq!( - children[0].name(), - "d", - "the child takes base's name for its id" - ); - assert_eq!(out.name(), "info", "the parent's name has not moved"); - } - - #[test] - fn a_column_the_generation_does_not_declare_is_left_alone() { - let m = mapping(vec![lance_field("value", 1, vec![])], &[(1, "amount")]); - let field = Field::new("_rowaddr", DataType::UInt64, true); - let mut renamed = false; - let out = m.rename(&field, &m.fields.clone(), &mut renamed); - assert_eq!(out.name(), "_rowaddr"); - assert!(!renamed); - } - fn create_test_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 2e3a2094b11..c2e98d15342 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -12,19 +12,18 @@ //! - [`IndexStore`] - In-memory index management //! - [`MemTableFlusher`] - Flush MemTable to storage as single Lance file -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::fmt::Debug; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, RwLock as StdRwLock}; use std::time::{Duration, Instant}; +use super::reconcile::{Plan, without_field_ids}; use arc_swap::ArcSwap; -use arrow::compute::CastOptions; use arrow_array::{ArrayRef, BooleanArray, RecordBatch, new_null_array}; -use arrow_cast::cast_with_options; -use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; +use arrow_schema::Schema as ArrowSchema; use async_trait::async_trait; -use lance_core::datatypes::{LANCE_FIELD_ID_KEY, Schema}; +use lance_core::datatypes::Schema; use lance_core::{Error, Result}; use lance_index::mem_wal::ShardManifest; use lance_index::vector::hnsw::builder::HnswBuildParams; @@ -1603,21 +1602,6 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// The lance field id an Arrow field carries, if it carries one. -/// -/// Lance already reads this key when converting Arrow to its own schema; the -/// memtable's storage schema carries it so that a WAL entry, which is Arrow IPC -/// and so keeps field metadata, stays addressable by id rather than by name -/// alone. Base data files have always been addressed this way -/// (`DataFile.fields` is a list of ids); this brings the fresh tier alongside. -fn field_id_of(field: &ArrowField) -> Option { - field - .metadata() - .get(LANCE_FIELD_ID_KEY) - .and_then(|v| v.parse::().ok()) - .filter(|id| *id >= 0) -} - /// Re-label `batch` to the storage schema, matching columns by **field id** /// where both sides carry one, and by **name** otherwise. /// @@ -1646,83 +1630,11 @@ fn conform_to_storage_schema( storage_schema: &Arc, pk_columns: &[String], ) -> Result { - let n = batch.num_rows(); - let by_field_id: HashMap = batch - .schema() - .fields() - .iter() - .enumerate() - .filter_map(|(i, f)| field_id_of(f).map(|id| (id, i))) - .collect(); - - // An entry column an id match has already taken. A rename frees a name for - // something else to use, so the entry's column can answer to the schema's - // name without being the schema's column -- and it is the id match that - // says which one it really is. - let claimed: HashSet = storage_schema - .fields() - .iter() - .filter_map(|f| field_id_of(f)) - .filter_map(|id| by_field_id.get(&id).copied()) - .collect(); - - let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); - for field in storage_schema.fields() { - let name = field.name(); - let carried = field_id_of(field) - .and_then(|id| by_field_id.get(&id).map(|i| batch.column(*i).clone())) - .or_else(|| { - // No column of this id. The entry may still carry this column - // under this name with an id of its own -- a cast takes a new - // id and keeps the name -- but only if nothing has claimed it. - let (i, _) = batch.schema().column_with_name(name)?; - (!claimed.contains(&i)).then(|| batch.column(i).clone()) - }); - if let Some(column) = carried { - columns.push(if column.data_type() == field.data_type() { - column - } else { - // `safe: false` so a lossy cast is an error rather than a - // column of nulls -- the same option `alter_columns` casts the - // base data under, so both halves of the table agree. - cast_with_options( - &column, - field.data_type(), - &CastOptions { - safe: false, - ..Default::default() - }, - ) - .map_err(|e| { - Error::invalid_input(format!( - "column '{name}' was written as {} and the schema now declares {}, \ - which it cannot be cast to: {e}", - column.data_type(), - field.data_type(), - )) - })? - }); - } else if name == TOMBSTONE { - columns.push(Arc::new(BooleanArray::from(vec![false; n]))); - } else if pk_columns.iter().any(|c| c == name) { - return Err(Error::invalid_input(format!( - "batch is missing primary key column '{}' declared by the storage schema", - name - ))); - } else { - // Non-primary-key columns are nullable in the storage schema - // whatever the base table declares (`relax_non_pk_nullability`), so - // a null stands in for a value the entry never held. - columns.push(new_null_array(field.data_type(), n)); - } + let plan = Plan::resolve(batch.schema().as_ref(), storage_schema, pk_columns)?; + if plan.is_identity() { + return Ok(batch); } - RecordBatch::try_new(storage_schema.clone(), columns).map_err(|e| { - Error::invalid_input(format!( - "failed to conform a batch to the storage schema \ - (does the batch match the base table schema?): {}", - e - )) - }) + plan.apply(&batch) } /// Build a tombstone batch from a key-only `keys` batch: primary keys carried @@ -2200,8 +2112,11 @@ impl ShardWriter { // The caller's schema is the shard's logical schema; the storage schema // is derived below, once the primary key is known. lance owns // `_tombstone` and appends it here — idempotent across reopens. - let logical_schema = schema; - let tombstoned = schema_with_tombstone(&logical_schema); + // The stored schema carries field ids so identity survives a rename; + // what a caller's batch is checked against must not, since a batch + // carries none and Arrow compares a struct's children in full. + let tombstoned = schema_with_tombstone(&schema); + let logical_schema = Arc::new(without_field_ids(&schema)); let base_uri = base_uri.into(); let shard_id = config.shard_id; diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 8c1977d1605..669dc848bad 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -12,6 +12,7 @@ use super::{ transaction::{Operation, Transaction}, write::cleanup_data_fragments, }; +use crate::dataset::mem_wal::DatasetMemWalExt; use crate::index::load_all_indices; use crate::{Error, Result, io::exec::Planner}; use arrow::compute::CastOptions; @@ -471,6 +472,29 @@ pub(super) async fn add_columns( .await } +/// Refuse to change a column's type on a table with a MemWAL attached. +/// +/// A cast gives the column a new field id and keeps its name, which is exactly +/// what dropping a column and adding another under that name looks like. Rows +/// the WAL still holds carry the old id, and nothing in the schemas says which +/// of the two happened -- so they could only be reconciled by guessing. +/// +/// A table without a MemWAL is unaffected: this is the only thing the check +/// looks at. +/// +/// Takes the decision as a bool rather than the alterations themselves: a +/// reference to them held across an await would have to be `Sync`. +async fn reject_cast_on_mem_wal(dataset: &Dataset, casts: bool) -> Result<()> { + if !casts || dataset.mem_wal_index_details().await?.is_none() { + return Ok(()); + } + Err(Error::invalid_input( + "cannot change a column's type on a table with a MemWAL attached: a cast takes a \ + new field id, which rows still in the WAL cannot be matched to. Drop the MemWAL, \ + or add a column of the new type and backfill it.", + )) +} + async fn cleanup_new_column_data_files(fragments: &[FileFragment], new_fragments: &[Fragment]) { let Some(first_fragment) = fragments.first() else { return; @@ -729,6 +753,8 @@ pub(super) async fn alter_columns( dataset: &mut Dataset, alterations: &[ColumnAlteration], ) -> Result<()> { + reject_cast_on_mem_wal(dataset, alterations.iter().any(|a| a.data_type.is_some())).await?; + // Validate referenced columns exist and enforce NOT NULL when tightening // a column from nullable to non-nullable. let mut new_schema = dataset.schema().clone(); From 37f41dd6466b6b14919718ec872ecf26a9ea0b72 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 16:09:53 -0400 Subject: [PATCH 19/48] fix(mem_wal): ask a generation for its own column names on the search paths Vector and full-text search open their own SSTable arms and projected the table's column names onto them. A rename moves the table's name while the file still holds the old one, so the projection asks for a column that is not there and the search fails. They translate the projection the same way the scan does, by field id. --- .../src/dataset/mem_wal/scanner/fts_search.rs | 19 +++++++++++++++-- .../src/dataset/mem_wal/scanner/planner.rs | 2 +- .../dataset/mem_wal/scanner/vector_search.rs | 21 +++++++++++++++++-- 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 5d8885d01cd..181c13c24bc 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -57,8 +57,10 @@ use super::block_list::compute_source_block_lists; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{FirstByPkExec, PkBlockFilterExec}; +use super::planner::stored_names; use super::projection::{project_to_canonical, validate_projection_names}; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; +use crate::dataset::mem_wal::arrow_schema_with_field_ids; use crate::dataset::mem_wal::memtable::scanner::MemTableScanner; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use crate::index::scalar::inverted::{ @@ -906,8 +908,21 @@ impl LsmFtsSearchPlanner { ) .await?; let mut scanner = dataset.scan(); - let cols = self.fts_scanner_projection(projection); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + // Asked of this generation under its own names: a rename moved + // the table's name while the file still holds the old one. + let stored = arrow_schema_with_field_ids(dataset.schema()); + let names = stored_names(&stored, &self.base_schema); + let wanted = self.fts_scanner_projection(projection); + let cols: Vec<&str> = wanted + .iter() + .filter_map(|name| { + names + .iter() + .find(|(_, table_name)| *table_name == name) + .map(|(stored_name, _)| stored_name.as_str()) + }) + .collect(); + scanner.project(&cols)?; if let Some(ref filter) = self.filter { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the BM25 top-k. diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 8964c3366a0..b654cbc9329 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -88,7 +88,7 @@ fn filter_above(plan: Arc, expr: &Expr) -> Result HashMap { +pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap { eprintln!( "DBG stored={:?} table={:?}", stored diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index c808a4e3a5a..399cb106142 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -28,11 +28,13 @@ use crate::io::exec::TakeExec; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; +use super::planner::stored_names; use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, validate_projection_names, wants_row_id, }; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; +use crate::dataset::mem_wal::arrow_schema_with_field_ids; use crate::session::Session; use lance_io::object_store::ObjectStoreParams; @@ -495,9 +497,24 @@ impl LsmVectorSearchPlanner { ) .await?; let mut scanner = dataset.scan(); - let cols = + // Asked of this generation under its own names: a rename moved + // the table's name while the file still holds the old one, so + // projecting the table's names would ask for a column that is + // not there. + let stored = arrow_schema_with_field_ids(dataset.schema()); + let names = stored_names(&stored, &self.base_schema); + let wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + let cols: Vec<&str> = wanted + .iter() + .filter_map(|name| { + names + .iter() + .find(|(_, table_name)| *table_name == name) + .map(|(stored_name, _)| stored_name.as_str()) + }) + .collect(); + scanner.project(&cols)?; if let Some(ref filter) = self.filter { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the top-k. From 384e0e5db17e5b40f4abdb6f4491b6a934f464d5 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 17:27:49 -0400 Subject: [PATCH 20/48] docs(mem_wal): describe the union's reconciliation as it now stands --- rust/lance/src/dataset/mem_wal/scanner/projection.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index dcab0d9157b..6e83dae547c 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs @@ -225,8 +225,9 @@ pub fn project_to_canonical( Some((idx, source)) if source.data_type() == field.data_type() => { Arc::new(Column::new(name, idx)) } - // A generation sealed before this column was retyped carries the - // type it was sealed under. + // Arms reaching the union are already reconciled to the table's + // types, so this is the base arm meeting a canonical schema that + // widens one -- a cast the table itself declares. Some((idx, _)) => Arc::new(CastExpr::new( Arc::new(Column::new(name, idx)), field.data_type().clone(), @@ -238,8 +239,8 @@ pub fn project_to_canonical( // its rows are all live. The column is non-nullable, so a null here // fails the scan outright. None if name == TOMBSTONE => Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))), - // A generation sealed before this column existed, or under the name - // it carried then. Typed nulls are what it holds for those rows. + // A source sealed before this column existed. Typed nulls are what + // it holds for those rows. None => Arc::new(Literal::new( ScalarValue::try_from(field.data_type()).map_err(|e| { lance_core::Error::internal(format!( From 2842ce18d97f18bc1a8cc92a5a9a24c4b991d189 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 17:36:24 -0400 Subject: [PATCH 21/48] refactor(mem_wal): drop what the shared plan made unnecessary The relabel exec no longer rebuilds struct arrays: arms are reconciled to the table's names before the union sees them, so the only relabel left is the logical/storage nullability boundary it was written for. Also corrects the comments that still described reconciling a retype. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 11 +++-- .../mem_wal/scanner/exec/schema_relabel.rs | 46 ++----------------- .../src/dataset/mem_wal/scanner/planner.rs | 8 ++-- 3 files changed, 14 insertions(+), 51 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 82dd2fd1bcd..f5423c88e95 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -9,9 +9,14 @@ //! it to a WAL entry, and a scan applies it to a generation's batches. //! //! Columns are matched by **field id**. A rename changes a field's name and -//! keeps its id, so a name is not identity; a cast keeps the name and takes a -//! new id, so a name is not identity there either. Where the two disagree the -//! plan refuses rather than guesses -- see [`Resolution::Ambiguous`]. +//! keeps its id, so a name is not identity: a source column of the same name +//! under a different id is a different column, and reading it would answer with +//! values the table no longer has. A name is matched only where identity is +//! absent, as in a batch a caller has just handed in. +//! +//! Nothing here has to tell a cast from a column dropped and replaced, which no +//! rule can: a table with a MemWAL refuses to change a column's type, so the +//! only way for an id to disappear is a drop. use std::collections::HashMap; use std::sync::Arc; diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index 1e525a24efc..88691f99d63 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -8,8 +8,8 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, StructArray}; -use arrow_schema::{DataType, SchemaRef}; +use arrow_array::{RecordBatch, RecordBatchOptions}; +use arrow_schema::SchemaRef; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::EquivalenceProperties; @@ -123,18 +123,9 @@ impl Stream for SchemaRelabelStream { Poll::Ready(Some(Ok(batch))) => { // Carry the row count explicitly: `try_new` infers it from the // first column, which a column-less batch does not have. - // A struct's child names live in the array's own type, not in - // the schema above it, so a renamed child needs the array - // rebuilt. The children are reused, so it costs a pointer copy. - let columns: Vec = batch - .columns() - .iter() - .zip(self.schema.fields()) - .map(|(column, field)| relabel_array(column, field.data_type())) - .collect(); let relabeled = RecordBatch::try_new_with_options( self.schema.clone(), - columns, + batch.columns().to_vec(), &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), ) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)); @@ -151,37 +142,6 @@ impl datafusion::physical_plan::RecordBatchStream for SchemaRelabelStream { } } -/// `array` with its type relabelled to `target`, rebuilding a struct whose -/// children are named differently. -/// -/// Only names differ here -- the layout is identical -- so an array whose type -/// already matches, and any type this cannot express, is returned as it is and -/// left for `RecordBatch::try_new` to reject. -fn relabel_array(array: &ArrayRef, target: &DataType) -> ArrayRef { - if array.data_type() == target { - return Arc::clone(array); - } - let (DataType::Struct(_), DataType::Struct(target_fields)) = (array.data_type(), target) else { - return Arc::clone(array); - }; - let Some(source) = array.as_any().downcast_ref::() else { - return Arc::clone(array); - }; - if source.columns().len() != target_fields.len() { - return Arc::clone(array); - } - let children: Vec = source - .columns() - .iter() - .zip(target_fields) - .map(|(child, field)| relabel_array(child, field.data_type())) - .collect(); - match StructArray::try_new(target_fields.clone(), children, source.nulls().cloned()) { - Ok(rebuilt) => Arc::new(rebuilt), - Err(_) => Arc::clone(array), - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index b654cbc9329..f5a36fb3486 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -107,8 +107,7 @@ pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap Date: Tue, 15 Sep 2026 17:42:02 -0400 Subject: [PATCH 22/48] test(mem_wal): drop the fixtures the removed mapping tests used --- .../src/dataset/mem_wal/scanner/planner.rs | 29 ------------------- 1 file changed, 29 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index f5a36fb3486..6adfc2ba9dd 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -630,35 +630,6 @@ mod tests { use super::*; use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; - /// A lance field, as a generation's schema records it. - fn lance_field(name: &str, id: i32, children: Vec) -> LanceField { - let arrow = if children.is_empty() { - Field::new(name, DataType::Int64, true) - } else { - Field::new( - name, - DataType::Struct( - children - .iter() - .map(|c| Field::new(&c.name, DataType::Int64, true)) - .collect(), - ), - true, - ) - }; - let mut field = LanceField::try_from(&arrow).expect("lance field"); - field.set_id(-1, &mut (id - 1).clone()); - field.id = id; - for (child, source) in field.children.iter_mut().zip(children.iter()) { - child.id = source.id; - } - field - } - - fn names(pairs: &[(i32, &str)]) -> HashMap { - pairs.iter().map(|(id, n)| (*id, n.to_string())).collect() - } - fn create_test_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), From 8c224760e3fcd82e0457c7b28f1d2e1e63639aa2 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 18:54:36 -0400 Subject: [PATCH 23/48] fix(mem_wal): make identity reach every nested field, and the batch match it A field id resolves which column is which, and a nested column carries its children inside its own type -- so the ids have to reach them and the arrays have to be built under them. Ids are stamped and stripped through a struct's children, a list's element and that element's children in turn, and restored the same way onto a scan's output, which is built from the dataset and carries none. A nested column's array is rebuilt under the target's children even when nothing about them moved: otherwise an unchanged struct produces a batch that disagrees with the schema it is built under, which Arrow rejects. Lists and fixed-size lists keep their offsets, width and validity; only the element's own type moves. A predicate over a rebuilt column is answered after reconciliation. Its reference names the parent -- a predicate on a struct's child refers to the struct -- and the parent's name does not move when a child is renamed, so a name check alone would push it down to be evaluated against the names this generation happens to hold. Live input is bound to its own schema before reconciliation. It is trusted for its values and not for its identity: a validated batch's columns are already the right ones in the right order, and any ids it carries are the caller's to claim rather than the table's to honour. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 146 ++++++++++++++++-- .../src/dataset/mem_wal/scanner/builder.rs | 11 +- .../src/dataset/mem_wal/scanner/fts_search.rs | 7 +- .../src/dataset/mem_wal/scanner/planner.rs | 62 ++++++-- .../dataset/mem_wal/scanner/vector_search.rs | 7 +- rust/lance/src/dataset/mem_wal/write.rs | 35 ++++- 6 files changed, 224 insertions(+), 44 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index f5423c88e95..9d2f6c609bb 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -21,7 +21,10 @@ use std::collections::HashMap; use std::sync::Arc; -use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions, StructArray}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, FixedSizeListArray, ListArray, RecordBatch, RecordBatchOptions, + StructArray, +}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::{Error, Result}; @@ -52,6 +55,18 @@ pub(crate) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { let children: Vec = children.iter().map(|c| strip(c)).collect(); field.with_data_type(DataType::Struct(children.into())) } + DataType::List(element) => field + .clone() + .with_data_type(DataType::List(Arc::new(strip(element)))), + DataType::LargeList(element) => field + .clone() + .with_data_type(DataType::LargeList(Arc::new(strip(element)))), + DataType::FixedSizeList(element, size) => { + let size = *size; + field + .clone() + .with_data_type(DataType::FixedSizeList(Arc::new(strip(element)), size)) + } _ => field, } } @@ -208,6 +223,17 @@ fn resolve_field( }; let source = &source_fields[index]; + // A nested column's children are part of the array's own type, metadata + // included, so the array is rebuilt under the target's children even when + // nothing about them moved -- otherwise the batch disagrees with the schema + // it is built under. The leaves are reused, so it costs a pointer copy. + if is_nested(field.data_type()) { + return Ok(Source::Nested( + index, + resolve_children(source, field)?, + field.data_type().clone(), + )); + } if source.data_type() == field.data_type() { return Ok(Source::Take(index)); } @@ -225,20 +251,54 @@ fn resolve_field( field.data_type() ))); } - // A struct whose children were renamed keeps the parent's name and id, so - // only the children differ. Resolve them the same way rather than casting, - // which Arrow matches by name. - if let (DataType::Struct(source_children), DataType::Struct(target_children)) = - (source.data_type(), field.data_type()) - { - let claimed = claimed_children(source_children, target_children); - let children = target_children - .iter() - .map(|child| resolve_field(child, source_children, &claimed, &[])) - .collect::>>()?; - return Ok(Source::Nested(index, children, field.data_type().clone())); + unreachable!("a struct is resolved above and any other mismatch is rejected") +} + +/// Whether this type carries its children inside its own type, so an array of +/// it has to be rebuilt rather than taken as it stands. +fn is_nested(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Struct(_) + | DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) + ) +} + +/// The child fields of a nested type, if it has them. +fn children_of(data_type: &DataType) -> Option { + match data_type { + DataType::Struct(children) => Some(children.clone()), + // A list has exactly one child: its element. Its name is part of the + // type, so it is resolved like any other. + DataType::List(element) | DataType::LargeList(element) => { + Some(vec![element.as_ref().clone()].into()) + } + DataType::FixedSizeList(element, _) => Some(vec![element.as_ref().clone()].into()), + _ => None, } - unreachable!("a non-struct type mismatch is rejected above") +} + +/// How each of `field`'s children is produced from `source`'s. +fn resolve_children(source: &ArrowField, field: &ArrowField) -> Result> { + let (Some(source_children), Some(target_children)) = ( + children_of(source.data_type()), + children_of(field.data_type()), + ) else { + return Err(Error::invalid_input(format!( + "column `{}` is stored as {} and the schema declares {}; a column's type \ + cannot change on a table with a MemWAL", + field.name(), + source.data_type(), + field.data_type() + ))); + }; + let claimed = claimed_children(&source_children, &target_children); + target_children + .iter() + .map(|child| resolve_field(child, &source_children, &claimed, &[])) + .collect() } fn claimed_children(source: &arrow_schema::Fields, target: &arrow_schema::Fields) -> Vec { @@ -259,9 +319,65 @@ fn claimed_children(source: &arrow_schema::Fields, target: &arrow_schema::Fields fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) -> Result { match source { Source::Take(i) => Ok(Arc::clone(&columns[*i])), + // A list is rebuilt around its element: the offsets and the validity + // say which rows hold what, and only the element's own type moves. + Source::Nested(i, children, to @ (DataType::List(_) | DataType::LargeList(_))) => { + let list = columns[*i] + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input(format!("column `{name}` is not a list")))?; + let Some(element) = children_of(to).and_then(|c| c.first().cloned()) else { + unreachable!("a list target has an element"); + }; + let child = take_column( + &children[0], + std::slice::from_ref(list.values()), + list.values().len(), + element.name(), + )?; + Ok(Arc::new( + ListArray::try_new( + element, + list.offsets().clone(), + child, + list.nulls().cloned(), + ) + .map_err(|e| Error::invalid_input(format!("rebuild list column `{name}`: {e}")))?, + )) + } + // A fixed-size list is rebuilt the same way, keeping its width. + Source::Nested(i, children, to @ DataType::FixedSizeList(_, _)) => { + let DataType::FixedSizeList(_, size) = to else { + unreachable!("matched above"); + }; + let list = columns[*i] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input(format!("column `{name}` is not a fixed-size list")) + })?; + let Some(element) = children_of(to).and_then(|c| c.first().cloned()) else { + unreachable!("a fixed-size list target has an element"); + }; + let child = take_column( + &children[0], + std::slice::from_ref(list.values()), + list.values().len(), + element.name(), + )?; + Ok(Arc::new( + FixedSizeListArray::try_new(element, *size, child, list.nulls().cloned()).map_err( + |e| { + Error::invalid_input(format!( + "rebuild fixed-size list column `{name}`: {e}" + )) + }, + )?, + )) + } Source::Nested(i, children, to) => { let DataType::Struct(target_children) = to else { - unreachable!("Nested is only built for a struct target"); + unreachable!("Nested is only built for a struct or list target"); }; let struct_array = columns[*i] .as_any() diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs index 8bd5ddb94bf..e7af05f250d 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs @@ -592,6 +592,7 @@ impl LsmScanner { collector, self.pk_columns.clone(), base_schema, + Arc::clone(&self.identity_schema), nearest.column.clone(), distance_type, ) @@ -660,9 +661,13 @@ impl LsmScanner { }; let collector = self.build_collector(); - let mut planner = - super::LsmFtsSearchPlanner::new(collector, self.pk_columns.clone(), base_schema) - .with_filter(self.filter.clone()); + let mut planner = super::LsmFtsSearchPlanner::new( + collector, + self.pk_columns.clone(), + base_schema, + Arc::clone(&self.identity_schema), + ) + .with_filter(self.filter.clone()); if let Some(session) = &self.session { planner = planner.with_session(session.clone()); } diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 181c13c24bc..e8682b33ff2 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -415,6 +415,9 @@ pub struct LsmFtsSearchPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, + /// The same schema with each field's id, which resolves a generation's + /// stored columns to the table's. + identity_schema: SchemaRef, /// Session threaded into SSTable opens (shared caches). session: Option>, /// Store params for opening SSTables, reusing the base dataset's store. @@ -437,11 +440,13 @@ impl LsmFtsSearchPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, + identity_schema: SchemaRef, ) -> Self { Self { collector, pk_columns, base_schema, + identity_schema, session: None, store_params: None, sstable_cache: None, @@ -911,7 +916,7 @@ impl LsmFtsSearchPlanner { // Asked of this generation under its own names: a rename moved // the table's name while the file still holds the old one. let stored = arrow_schema_with_field_ids(dataset.schema()); - let names = stored_names(&stored, &self.base_schema); + let names = stored_names(&stored, &self.identity_schema); let wanted = self.fts_scanner_projection(projection); let cols: Vec<&str> = wanted .iter() diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 6adfc2ba9dd..6442d43f7b0 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -82,6 +82,15 @@ fn filter_above(plan: Arc, expr: &Expr) -> Result bool { + matches!(data_type, DataType::Struct(_)) +} + /// What the table calls each of this generation's stored columns, keyed by the /// stored name. /// @@ -89,19 +98,6 @@ fn filter_above(plan: Arc, expr: &Expr) -> Result HashMap { - eprintln!( - "DBG stored={:?} table={:?}", - stored - .fields() - .iter() - .map(|f| (f.name().clone(), field_id_of(f))) - .collect::>(), - table - .fields() - .iter() - .map(|f| (f.name().clone(), field_id_of(f))) - .collect::>() - ); let by_id: HashMap = table .fields() .iter() @@ -147,7 +143,8 @@ fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { let field = field.clone().with_metadata(metadata); // A struct's children carry their own ids, and a child can be renamed // while its parent's name does not move. - match (field.data_type(), source.data_type()) { + let data_type = field.data_type().clone(); + match (&data_type, source.data_type()) { (DataType::Struct(children), DataType::Struct(source_children)) => { let children: Vec = children .iter() @@ -155,6 +152,27 @@ fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { .collect(); field.with_data_type(DataType::Struct(children.into())) } + // A list's element is a field with an id of its own, and so are its + // children in turn. + (DataType::List(element), DataType::List(source_element)) => { + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::List(Arc::new(restore(element, &one)))) + } + (DataType::LargeList(element), DataType::LargeList(source_element)) => { + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::LargeList(Arc::new(restore(element, &one)))) + } + ( + DataType::FixedSizeList(element, size), + DataType::FixedSizeList(source_element, _), + ) => { + let size = *size; + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::FixedSizeList( + Arc::new(restore(element, &one)), + size, + )) + } _ => field, } } @@ -512,10 +530,20 @@ impl LsmScanPlanner { // A predicate this generation cannot answer as written runs // after reconciliation, reading its columns from this scan -- // so they have to be in it whether the caller asked or not. + // A predicate can run here only against columns this generation + // stores under the table's own name and shape. A nested column + // names its parent in the reference -- `info.a` refers to + // `info` -- and the parent's name does not move when a child is + // renamed, so a name check alone would push the predicate down + // to be evaluated against the child names this generation has. + // Anything reconstructed is answered after reconciliation. let answerable = filter.is_none_or(|expr| { - expr.column_refs() - .iter() - .all(|c| names.get(c.name.as_str()) == Some(&c.name)) + expr.column_refs().iter().all(|c| { + names.get(c.name.as_str()) == Some(&c.name) + && stored + .field_with_name(&c.name) + .is_ok_and(|f| !is_reconstructed(f.data_type())) + }) }); if let Some(expr) = filter.filter(|_| !answerable) { for column in expr.column_refs() { diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 399cb106142..d258f299d7c 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -83,6 +83,9 @@ pub struct LsmVectorSearchPlanner { pk_columns: Vec, /// Schema of the base table. base_schema: SchemaRef, + /// The same schema with each field's id, which resolves a generation's + /// stored columns to the table's. + identity_schema: SchemaRef, /// Vector column name. vector_column: String, /// Distance metric type (L2, Cosine, Dot, etc.). @@ -129,6 +132,7 @@ impl LsmVectorSearchPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, + identity_schema: SchemaRef, vector_column: String, distance_type: lance_linalg::distance::DistanceType, ) -> Self { @@ -136,6 +140,7 @@ impl LsmVectorSearchPlanner { collector, pk_columns, base_schema, + identity_schema, vector_column, distance_type, dataset: None, @@ -502,7 +507,7 @@ impl LsmVectorSearchPlanner { // projecting the table's names would ask for a column that is // not there. let stored = arrow_schema_with_field_ids(dataset.schema()); - let names = stored_names(&stored, &self.base_schema); + let names = stored_names(&stored, &self.identity_schema); let wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); let cols: Vec<&str> = wanted diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index c2e98d15342..4a9b8e61a40 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -20,7 +20,7 @@ use std::time::{Duration, Instant}; use super::reconcile::{Plan, without_field_ids}; use arc_swap::ArcSwap; -use arrow_array::{ArrayRef, BooleanArray, RecordBatch, new_null_array}; +use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions, new_null_array}; use arrow_schema::Schema as ArrowSchema; use async_trait::async_trait; use lance_core::datatypes::Schema; @@ -1625,6 +1625,31 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, /// /// A primary key the batch does not carry stays an error — there is no value to /// invent. +/// A batch a caller just handed in, under the storage schema. +/// +/// Live input is trusted for its values and not for its identity: it has been +/// validated against the logical schema, so its columns are the right ones in +/// the right order, and any field ids it carries are the caller's to claim +/// rather than the table's to honour. Dropping them leaves the columns matched +/// by name, which is what a validated batch's order already established. +/// +/// A replayed entry is the opposite: its ids are the table's, written when the +/// entry was, and they are the only thing that survives a rename. +fn conform_live_batch( + batch: RecordBatch, + storage_schema: &Arc, + pk_columns: &[String], +) -> Result { + let plain = Arc::new(without_field_ids(batch.schema().as_ref())); + let batch = RecordBatch::try_new_with_options( + plain, + batch.columns().to_vec(), + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| Error::invalid_input(format!("bind a live batch to its own schema: {e}")))?; + conform_to_storage_schema(batch, storage_schema, pk_columns) +} + fn conform_to_storage_schema( batch: RecordBatch, storage_schema: &Arc, @@ -2665,9 +2690,7 @@ impl ShardWriter { // `_tombstone`. let batches = batches .into_iter() - .map(|b| { - conform_to_storage_schema(b, &writer_state.schema, &writer_state.pk_columns) - }) + .map(|b| conform_live_batch(b, &writer_state.schema, &writer_state.pk_columns)) .collect::>>()?; self.put_memtable(batches, state, writer_state, backpressure) .await @@ -2792,9 +2815,7 @@ impl ShardWriter { // Mirrors `put`. let batches = batches .into_iter() - .map(|b| { - conform_to_storage_schema(b, &writer_state.schema, &writer_state.pk_columns) - }) + .map(|b| conform_live_batch(b, &writer_state.schema, &writer_state.pk_columns)) .collect::>>()?; self.put_memtable_no_wait(batches, state, writer_state, backpressure) .await From 5438c1578597cba7d92e0bb4b51b3ce07a94812d Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 20:47:31 -0400 Subject: [PATCH 24/48] fix(mem_wal): read every sealed generation through one resolution A sealed generation is written under whatever the table's schema was at the time, so the names it stores are its own. Only the scan resolved them to the table's; the point lookup projected the table's names straight at the file, and the vector and full-text searches translated the projection but not the predicate, the searched column, or the output. GenerationRead is that resolution, and all four read paths now go through it: what to project, whether a predicate can be pushed down and under which names, and how to bring the result back. Three consequences the paths had each got differently: - a generation's own columns (_tombstone, anything the table has since dropped) are numbered in its own schema, so their ids collide with whatever the table gave those numbers. They are stripped before resolution rather than resolved to a table column that shares an id. - a reader is handed the table's plain schema, so reconciliation emits it. Field ids live inside a nested column's own type, and only replay, which writes back into the memtable's storage schema, keeps them. - a search whose column the generation predates has nothing to offer, and contributes an empty arm rather than failing the whole query. SchemaRelabelExec relabels the arrays as well as the schema, for the same reason: a nested column carries its ids in its type, so relabelling the schema alone left the two disagreeing. LsmScanner and LsmPointLookupPlanner take the id-carrying schema through with_identity_schema, which LsmScanner::new fills in from the dataset. A caller that supplies none matches by name, as before ids existed. --- rust/lance/src/dataset/mem_wal.rs | 36 +- rust/lance/src/dataset/mem_wal/reconcile.rs | 124 +++++-- rust/lance/src/dataset/mem_wal/scanner.rs | 1 + .../src/dataset/mem_wal/scanner/builder.rs | 33 +- .../mem_wal/scanner/exec/schema_relabel.rs | 29 +- .../src/dataset/mem_wal/scanner/fts_search.rs | 70 ++-- .../src/dataset/mem_wal/scanner/generation.rs | 329 ++++++++++++++++++ .../src/dataset/mem_wal/scanner/planner.rs | 214 +----------- .../dataset/mem_wal/scanner/point_lookup.rs | 57 ++- .../dataset/mem_wal/scanner/vector_search.rs | 66 ++-- rust/lance/src/dataset/mem_wal/write.rs | 24 +- 11 files changed, 665 insertions(+), 318 deletions(-) create mode 100644 rust/lance/src/dataset/mem_wal/scanner/generation.rs diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index ed376b7a5bf..d9c88802e35 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -135,7 +135,7 @@ pub fn relax_non_pk_nullability( /// Scoped to the memtable path deliberately: emitting the id from the global /// Arrow conversion would change every schema Lance hands out, including for /// callers that compare schemas for equality. -pub(crate) fn arrow_schema_with_field_ids(schema: &Schema) -> ArrowSchema { +pub fn arrow_schema_with_field_ids(schema: &Schema) -> ArrowSchema { let arrow: ArrowSchema = schema.into(); let fields: Vec = arrow .fields() @@ -162,14 +162,32 @@ fn stamp_field_id(field: &ArrowField, among: &[Field]) -> ArrowField { } _ => field.clone(), }; - let DataType::Struct(children) = field.data_type() else { - return field; - }; - let children: Vec = children - .iter() - .map(|child| stamp_field_id(child, &source.children)) - .collect(); - field.with_data_type(DataType::Struct(children.into())) + // A container carries its children inside its own type, and each of them is + // a field with an id of its own: a list's element, and that element's + // children in turn. + match field.data_type() { + DataType::Struct(children) => { + let children: Vec = children + .iter() + .map(|child| stamp_field_id(child, &source.children)) + .collect(); + field.with_data_type(DataType::Struct(children.into())) + } + DataType::List(element) => { + let element = stamp_field_id(element, &source.children); + field.with_data_type(DataType::List(Arc::new(element))) + } + DataType::LargeList(element) => { + let element = stamp_field_id(element, &source.children); + field.with_data_type(DataType::LargeList(Arc::new(element))) + } + DataType::FixedSizeList(element, size) => { + let size = *size; + let element = stamp_field_id(element, &source.children); + field.with_data_type(DataType::FixedSizeList(Arc::new(element), size)) + } + _ => field, + } } pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 9d2f6c609bb..41f60e4ce8e 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -22,8 +22,8 @@ use std::collections::HashMap; use std::sync::Arc; use arrow_array::{ - Array, ArrayRef, BooleanArray, FixedSizeListArray, ListArray, RecordBatch, RecordBatchOptions, - StructArray, + Array, ArrayRef, BooleanArray, FixedSizeListArray, GenericListArray, RecordBatch, + RecordBatchOptions, StructArray, }; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; use lance_core::datatypes::LANCE_FIELD_ID_KEY; @@ -45,6 +45,37 @@ pub(crate) fn field_id_of(field: &ArrowField) -> Option { /// Ids belong to the stored schema, where identity has to survive a rename. A /// caller's batch carries none, and Arrow compares a struct's children by their /// full field -- metadata included -- so a stamped schema would reject it. +/// One stored column under the type the caller declared. +/// +/// A field id lives inside a nested column's own Arrow type, so the memtable's +/// id-carrying storage schema and the table's plain one describe the same +/// values as two different types. Only the labels differ, so this relabels the +/// array rather than converting it. +pub(crate) fn relabel_to(column: &ArrayRef, data_type: &DataType) -> Result { + if column.data_type() == data_type { + return Ok(column.clone()); + } + let relabelled = column + .to_data() + .into_builder() + .data_type(data_type.clone()) + .build() + .map_err(|e| { + Error::invalid_input(format!( + "a {} column cannot be read as {data_type}: {e}", + column.data_type() + )) + })?; + Ok(arrow_array::make_array(relabelled)) +} + +/// [`without_field_ids`] for a type rather than a schema, for the nested types a +/// reconciliation builds. +fn without_field_ids_in(data_type: &DataType) -> DataType { + let one = ArrowSchema::new(vec![ArrowField::new("", data_type.clone(), true)]); + without_field_ids(&one).field(0).data_type().clone() +} + pub(crate) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { fn strip(field: &ArrowField) -> ArrowField { let mut metadata = field.metadata().clone(); @@ -99,6 +130,30 @@ pub struct Plan { } impl Plan { + /// The same plan emitting the table's plain Arrow schema. + /// + /// Resolution needs field ids on the target -- a rename keeps the id and + /// moves the name -- but a reader is handed the table's schema, which does + /// not carry them. They live inside a nested column's own type, so a struct + /// built to the id-carrying target is a different Arrow type from the one + /// the caller declared. Only replay, which writes back into the memtable's + /// id-carrying storage schema, keeps them. + pub(crate) fn emitting_plain_schema(mut self) -> Self { + fn strip(source: &mut Source) { + match source { + Source::Nested(_, children, data_type) => { + *data_type = without_field_ids_in(data_type); + children.iter_mut().for_each(strip); + } + Source::Null(data_type) => *data_type = without_field_ids_in(data_type), + _ => {} + } + } + self.sources.iter_mut().for_each(strip); + self.target = Arc::new(without_field_ids(&self.target)); + self + } + /// Resolve `source` against `target`, or say why it cannot be done. /// /// `pk_columns` may not be filled with nulls: a row with no primary key @@ -316,34 +371,49 @@ fn claimed_children(source: &arrow_schema::Fields, target: &arrow_schema::Fields claimed } +/// One list column rebuilt around a reconciled element, for either offset width. +fn rebuild_list( + column: &ArrayRef, + element_source: &Source, + target: &DataType, + name: &str, +) -> Result { + let list = column + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::invalid_input(format!("column `{name}` is not a list")))?; + let Some(element) = children_of(target).and_then(|c| c.first().cloned()) else { + unreachable!("a list target has an element"); + }; + let child = take_column( + element_source, + std::slice::from_ref(list.values()), + list.values().len(), + element.name(), + )?; + Ok(Arc::new( + GenericListArray::::try_new( + element, + list.offsets().clone(), + child, + list.nulls().cloned(), + ) + .map_err(|e| Error::invalid_input(format!("rebuild list column `{name}`: {e}")))?, + )) +} + fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) -> Result { match source { Source::Take(i) => Ok(Arc::clone(&columns[*i])), - // A list is rebuilt around its element: the offsets and the validity - // say which rows hold what, and only the element's own type moves. - Source::Nested(i, children, to @ (DataType::List(_) | DataType::LargeList(_))) => { - let list = columns[*i] - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::invalid_input(format!("column `{name}` is not a list")))?; - let Some(element) = children_of(to).and_then(|c| c.first().cloned()) else { - unreachable!("a list target has an element"); - }; - let child = take_column( - &children[0], - std::slice::from_ref(list.values()), - list.values().len(), - element.name(), - )?; - Ok(Arc::new( - ListArray::try_new( - element, - list.offsets().clone(), - child, - list.nulls().cloned(), - ) - .map_err(|e| Error::invalid_input(format!("rebuild list column `{name}`: {e}")))?, - )) + // A list is rebuilt around its element: the offsets and the validity say + // which rows hold what, and only the element's own type moves. The two + // offset widths are different array types and neither downcasts to the + // other. + Source::Nested(i, children, to @ DataType::List(_)) => { + rebuild_list::(&columns[*i], &children[0], to, name) + } + Source::Nested(i, children, to @ DataType::LargeList(_)) => { + rebuild_list::(&columns[*i], &children[0], to, name) } // A fixed-size list is rebuilt the same way, keeping its width. Source::Nested(i, children, to @ DataType::FixedSizeList(_, _)) => { diff --git a/rust/lance/src/dataset/mem_wal/scanner.rs b/rust/lance/src/dataset/mem_wal/scanner.rs index 907250a1cb1..fa73d727811 100644 --- a/rust/lance/src/dataset/mem_wal/scanner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner.rs @@ -42,6 +42,7 @@ mod collector; mod data_source; pub mod exec; mod fts_search; +mod generation; mod planner; mod point_lookup; mod projection; diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs index e7af05f250d..6794c97b325 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs @@ -308,9 +308,8 @@ impl LsmScanner { pk_columns: Vec, ) -> Self { Self { - // Whatever identity the caller supplied. Pass a schema carrying - // field ids to have a generation's columns resolved by them; a - // plain one is matched by name, as it was before ids existed. + // Name matching until the caller supplies ids, as it was before + // ids existed. See [`Self::with_identity_schema`]. identity_schema: schema.clone(), base: BaseSource::PathOnly(base_path.into()), schema, @@ -353,6 +352,18 @@ impl LsmScanner { /// flush) captured atomically by `ShardWriter::in_memory_memtable_refs`. /// The read path's entry point — closes the concurrent-read-vs-flush /// hole by carrying frozen-undrained generations into the scan. + /// Supply `schema` with each field's id, so a sealed generation's columns + /// are resolved to the table's by id rather than by name. A rename keeps + /// the id and moves the name, so without this a renamed column reads as + /// absent. Built with + /// [`arrow_schema_with_field_ids`](crate::dataset::mem_wal::arrow_schema_with_field_ids). + /// + /// Set for you by [`Self::new`], which has the dataset to read them from. + pub fn with_identity_schema(mut self, identity_schema: SchemaRef) -> Self { + self.identity_schema = identity_schema; + self + } + pub fn with_in_memory_memtables( mut self, shard_id: Uuid, @@ -592,10 +603,10 @@ impl LsmScanner { collector, self.pk_columns.clone(), base_schema, - Arc::clone(&self.identity_schema), nearest.column.clone(), distance_type, ) + .with_identity_schema(Arc::clone(&self.identity_schema)) .with_filter(self.filter.clone()); if let BaseSource::Table(dataset) = &self.base { planner = planner.with_dataset(dataset.clone()); @@ -661,13 +672,10 @@ impl LsmScanner { }; let collector = self.build_collector(); - let mut planner = super::LsmFtsSearchPlanner::new( - collector, - self.pk_columns.clone(), - base_schema, - Arc::clone(&self.identity_schema), - ) - .with_filter(self.filter.clone()); + let mut planner = + super::LsmFtsSearchPlanner::new(collector, self.pk_columns.clone(), base_schema) + .with_identity_schema(Arc::clone(&self.identity_schema)) + .with_filter(self.filter.clone()); if let Some(session) = &self.session { planner = planner.with_session(session.clone()); } @@ -712,7 +720,8 @@ impl LsmScanner { extract_pk_point_keys(filter, &self.pk_columns[0], pk_field.data_type()) { let mut planner = - LsmPointLookupPlanner::new(collector, self.pk_columns.clone(), base_schema); + LsmPointLookupPlanner::new(collector, self.pk_columns.clone(), base_schema) + .with_identity_schema(Arc::clone(&self.identity_schema)); if let Some(session) = &self.session { planner = planner.with_session(session.clone()); } diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index 88691f99d63..ddf1d017d73 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -18,6 +18,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use futures::{Stream, StreamExt}; +use crate::dataset::mem_wal::reconcile::relabel_to; /// Re-labels every batch to an exact target schema, leaving the arrays /// untouched. `ProjectionExec` cannot: DataFusion derives output nullability @@ -123,12 +124,26 @@ impl Stream for SchemaRelabelStream { Poll::Ready(Some(Ok(batch))) => { // Carry the row count explicitly: `try_new` infers it from the // first column, which a column-less batch does not have. - let relabeled = RecordBatch::try_new_with_options( - self.schema.clone(), - batch.columns().to_vec(), - &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), - ) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)); + // A nested column carries its field ids inside its own type, so + // relabelling the schema alone would leave the arrays + // disagreeing with it. + let columns: Result, _> = batch + .columns() + .iter() + .zip(self.schema.fields()) + .map(|(column, field)| { + relabel_to(column, field.data_type()) + .map_err(|e| DataFusionError::External(Box::new(e))) + }) + .collect(); + let relabeled = columns.and_then(|columns| { + RecordBatch::try_new_with_options( + self.schema.clone(), + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) + }); Poll::Ready(Some(relabeled)) } other => other, @@ -238,7 +253,7 @@ mod tests { let error = run(relabeled).await.unwrap_err().to_string(); assert!( - error.contains("column types must match"), + error.contains("cannot be read as Int32"), "expected a data type error, got: {error}" ); } diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index e8682b33ff2..f995c2c2757 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -57,10 +57,9 @@ use super::block_list::compute_source_block_lists; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{FirstByPkExec, PkBlockFilterExec}; -use super::planner::stored_names; +use super::generation::{GenerationRead, filter_above}; use super::projection::{project_to_canonical, validate_projection_names}; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; -use crate::dataset::mem_wal::arrow_schema_with_field_ids; use crate::dataset::mem_wal::memtable::scanner::MemTableScanner; use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; use crate::index::scalar::inverted::{ @@ -440,13 +439,12 @@ impl LsmFtsSearchPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, - identity_schema: SchemaRef, ) -> Self { Self { collector, pk_columns, + identity_schema: base_schema.clone(), base_schema, - identity_schema, session: None, store_params: None, sstable_cache: None, @@ -473,6 +471,16 @@ impl LsmFtsSearchPlanner { } /// Set the session used to open SSTables. + /// The table's schema carrying each field's id, which is what resolves a + /// generation's stored columns to the table's. + /// + /// Defaults to the base schema, so a caller that has no ids to give is + /// matched by name as it was. + pub fn with_identity_schema(mut self, schema: SchemaRef) -> Self { + self.identity_schema = schema; + self + } + pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); self @@ -789,6 +797,7 @@ impl LsmFtsSearchPlanner { *fetch_limit, projection, index_params.as_ref(), + &target_schema, )) })) .await?; @@ -881,6 +890,9 @@ impl LsmFtsSearchPlanner { limit: Option, projection: Option<&[String]>, index_params: Option<&InvertedIndexParams>, + // What every arm is normalized to, so an arm with nothing to offer can + // stand in for itself. + target_schema: &SchemaRef, ) -> Result> { match source { LsmDataSource::BaseTable { dataset } => { @@ -915,33 +927,51 @@ impl LsmFtsSearchPlanner { let mut scanner = dataset.scan(); // Asked of this generation under its own names: a rename moved // the table's name while the file still holds the old one. - let stored = arrow_schema_with_field_ids(dataset.schema()); - let names = stored_names(&stored, &self.identity_schema); let wanted = self.fts_scanner_projection(projection); - let cols: Vec<&str> = wanted - .iter() - .filter_map(|name| { - names - .iter() - .find(|(_, table_name)| *table_name == name) - .map(|(stored_name, _)| stored_name.as_str()) - }) - .collect(); - scanner.project(&cols)?; - if let Some(ref filter) = self.filter { + let mut generation = GenerationRead::new( + dataset.schema(), + Arc::clone(&self.identity_schema), + self.pk_columns.clone(), + wanted, + ); + // The index is on this generation's own column, under the name + // it had when the generation was sealed. + let Some(stored_column) = generation.stored_name(column) else { + // The generation was sealed before the searched column + // existed, so it has nothing to match. + return self.empty_plan(target_schema); + }; + let stored_column = stored_column.to_string(); + // A predicate this generation cannot answer as written runs + // above the reconciliation, where the columns it names exist. + // The BM25 top-k has already run by then, so the arm can come + // back short -- which is right: those rows have no value for a + // column sealed before it existed. + let pushed = self.filter.as_ref().map(|f| (f, generation.to_stored(f))); + if let Some((expr, None)) = pushed { + for column in expr.column_refs() { + generation.also_produce(&column.name); + } + } + scanner.project(&generation.stored_projection())?; + if let Some((_, Some(ref stored))) = pushed { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the BM25 top-k. - scanner.filter_expr(filter.clone()); + scanner.filter_expr(stored.clone()); scanner.prefilter(true); } - let mut bound_query = query.clone().with_column(column.to_string())?; + let mut bound_query = query.clone().with_column(stored_column)?; if let Some(limit) = limit { bound_query = bound_query.limit(Some(limit as i64)); } else { bound_query = bound_query.limit(None); } scanner.full_text_search(bound_query)?; - scanner.create_plan().await + let reconciled = generation.reconcile(scanner.create_plan().await?)?; + match pushed { + Some((expr, None)) => filter_above(reconciled, expr), + _ => Ok(reconciled), + } } LsmDataSource::ActiveMemTable { batch_store, diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs new file mode 100644 index 00000000000..fd4a9940be7 --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Reading a sealed generation under the table's current schema. +//! +//! A sealed generation is a Lance dataset of its own, written under whatever +//! the table's schema was at the time. The names it stores are therefore its +//! own: a rename since the seal moved the table's name and left the file +//! holding the old one, and only field ids relate the two. Every read path — +//! scan, point lookup, vector search, full-text search — goes through +//! [`GenerationRead`] so they resolve a generation the same way. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion::common::DFSchema; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::context::ExecutionProps; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::prelude::Expr; +use datafusion_physical_expr::create_physical_expr; +use lance_core::is_system_column; +use lance_core::{Error, Result}; + +use super::exec::ReconcileExec; +use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; +use lance_core::datatypes::LANCE_FIELD_ID_KEY; +use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; + +/// One sealed generation, read under the table's schema. +/// +/// Built from the generation's own schema and the table's id-carrying schema. +/// It answers three questions, in the order a scan needs them: what to project +/// from the file ([`Self::stored_projection`]), whether a predicate can be +/// pushed into it ([`Self::can_answer`] / [`Self::to_stored`]), and how to +/// bring the result back to the table's names and shapes +/// ([`Self::reconcile`]). +pub(super) struct GenerationRead { + /// The generation's own schema, carrying its field ids. + stored: Schema, + /// The generation's name for a column → the table's name for it. Columns + /// the table no longer has are absent. + names: HashMap, + /// The table's schema, carrying field ids. + identity: SchemaRef, + pk_columns: Vec, + /// Table names this read produces, in order. + wanted: Vec, +} + +impl GenerationRead { + /// `wanted` is what the caller asks for, in the table's names. + pub(super) fn new( + dataset_schema: &lance_core::datatypes::Schema, + identity: SchemaRef, + pk_columns: Vec, + wanted: Vec, + ) -> Self { + let stored = arrow_schema_with_field_ids(dataset_schema); + let names = stored_names(&stored, &identity); + Self { + stored, + names, + identity, + pk_columns, + wanted, + } + } + + /// The generation's name for one of the table's columns. + pub(super) fn stored_name(&self, table_name: &str) -> Option<&str> { + self.names + .iter() + .find(|(_, table)| *table == table_name) + .map(|(stored, _)| stored.as_str()) + } + + /// Also produce `column`, which the caller needs even though it did not ask + /// for it — a predicate that runs after reconciliation reads its columns + /// from this scan. + pub(super) fn also_produce(&mut self, column: &str) { + if !self.wanted.iter().any(|w| w == column) { + self.wanted.push(column.to_string()); + } + } + + /// What to project from the file: the wanted columns under the names it has. + /// A column it never stored is dropped here and filled in by + /// [`Self::reconcile`]. + pub(super) fn stored_projection(&self) -> Vec<&str> { + self.wanted + .iter() + .filter_map(|name| self.stored_name(name).or_else(|| self.stored_system_column(name))) + .collect() + } + + /// A system column (`_tombstone`, `_rowaddr`) is not one of the table's, so + /// no field id relates it; the generation stores it under the name it is + /// asked for, or not at all. + fn stored_system_column(&self, name: &str) -> Option<&str> { + (is_system_column(name) || name == TOMBSTONE) + .then(|| self.stored.field_with_name(name).ok()) + .flatten() + .map(|f| f.name().as_str()) + } + + /// Whether `expr` can be pushed into this generation's scan. + /// + /// It can when every column it names is stored under the table's own name + /// and shape. A nested column fails the shape half: a reference names the + /// parent (`info.a` refers to `info`) and a parent's name does not move + /// when a child is renamed, so the pushed-down predicate would be evaluated + /// against child names this generation has and the table does not. + pub(super) fn can_answer(&self, expr: &Expr) -> bool { + expr.column_refs().iter().all(|c| { + self.names.get(c.name.as_str()) == Some(&c.name) + && self + .stored + .field_with_name(&c.name) + .is_ok_and(|f| !is_reconstructed(f.data_type())) + }) + } + + /// `expr` with each column reference moved to the name this generation + /// stores it under. `None` when a column it names is not stored here at + /// all, or is reconstructed — then the predicate belongs above the scan. + pub(super) fn to_stored(&self, expr: &Expr) -> Option { + let pushable = expr.column_refs().iter().all(|c| { + self.stored_name(&c.name).is_some_and(|stored| { + self.stored + .field_with_name(stored) + .is_ok_and(|f| !is_reconstructed(f.data_type())) + }) + }); + if !pushable { + return None; + } + expr.clone() + .transform(|e| match e { + Expr::Column(mut c) => { + // `stored_name` is total over the refs, checked above. + let stored = self.stored_name(&c.name).expect("checked").to_string(); + c.name = stored; + Ok(Transformed::yes(Expr::Column(c))) + } + other => Ok(Transformed::no(other)), + }) + .map(|t| t.data) + .ok() + } + + /// Bring the scan's output back to the table's names and shapes: renames + /// followed, columns the generation never stored filled with nulls, nested + /// columns rebuilt to the shape the table declares. + /// + /// `scan` may produce more than was asked for (`_rowaddr`, `_tombstone`); + /// those pass through untouched, as does anything the generation has that + /// the table does not. + pub(super) fn reconcile(&self, scan: Arc) -> Result> { + let source = self.only_the_tables_ids(with_ids_from(&scan.schema(), &self.stored)); + let target = self.target(&source); + let plan = Plan::resolve(&source, &target, &self.pk_columns)?.emitting_plain_schema(); + if plan.is_identity() { + return Ok(scan); + } + Ok(Arc::new(ReconcileExec::new(scan, Arc::new(plan)))) + } + + /// `source` with the field ids of everything that is not one of the + /// table's columns removed. + /// + /// A generation numbers its own columns in its own schema -- `_tombstone`, + /// and anything it holds that the table has since dropped -- so those ids + /// collide with whatever the table gave those numbers. Left in place, a + /// column added to the table resolves to whichever of them happens to share + /// its id. + fn only_the_tables_ids(&self, source: Schema) -> Schema { + let fields: Vec = source + .fields() + .iter() + .map(|field| match self.names.contains_key(field.name()) { + true => field.as_ref().clone(), + false => { + let mut metadata = field.metadata().clone(); + metadata.remove(LANCE_FIELD_ID_KEY); + field.as_ref().clone().with_metadata(metadata) + } + }) + .collect(); + Schema::new_with_metadata(fields, source.metadata().clone()) + } + + /// The schema [`Self::reconcile`] produces: the wanted columns as the table + /// declares them, then whatever else the scan carries. + fn target(&self, source: &Schema) -> SchemaRef { + let mut fields: Vec = self + .wanted + .iter() + .filter_map(|name| self.identity.field_with_name(name).ok().cloned()) + .collect(); + // A generation's own columns are not the table's, so they pass through + // as the generation has them. + for field in source.fields() { + let is_the_tables = self.names.contains_key(field.name()); + if !is_the_tables && fields.iter().all(|f| f.name() != field.name()) { + fields.push(field.as_ref().clone()); + } + } + Arc::new(Schema::new(fields)) + } +} + +/// Run `expr` above `plan`, for a predicate the generation could not answer +/// as written. +pub(super) fn filter_above( + plan: Arc, + expr: &Expr, +) -> Result> { + let schema = plan.schema(); + let df_schema = DFSchema::try_from(schema.as_ref().clone()) + .map_err(|e| Error::internal(format!("filter schema: {e}")))?; + let props = ExecutionProps::new(); + let physical = create_physical_expr(expr, &df_schema, &props) + .map_err(|e| Error::internal(format!("plan filter `{expr}`: {e}")))?; + Ok(Arc::new( + FilterExec::try_new(physical, plan).map_err(|e| Error::internal(format!("filter: {e}")))?, + )) +} + +/// A column the reconciliation rebuilds rather than takes as it stands, so its +/// stored shape is not the shape a pushed-down predicate would expect. +fn is_reconstructed(data_type: &DataType) -> bool { + match data_type { + DataType::Struct(_) => true, + DataType::List(e) | DataType::LargeList(e) | DataType::FixedSizeList(e, _) => { + is_reconstructed(e.data_type()) + } + _ => false, + } +} + +/// The generation's name for each of the table's columns, by field id: a rename +/// changes the name and keeps the id. +pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap { + let by_id: HashMap = table + .fields() + .iter() + .filter_map(|f| field_id_of(f).map(|id| (id, f.name().as_str()))) + .collect(); + // A caller that supplies no ids leaves only names to match on. + if by_id.is_empty() { + return stored + .fields() + .iter() + .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) + .filter(|f| table.field_with_name(f.name()).is_ok()) + .map(|f| (f.name().clone(), f.name().clone())) + .collect(); + } + stored + .fields() + .iter() + // A generation's own columns are numbered in its own schema, so their + // ids collide with whatever the table gave those numbers. They are not + // the table's columns and are never resolved to one. + .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) + .filter_map(|f| { + let id = field_id_of(f)?; + by_id + .get(&id) + .map(|name| (f.name().clone(), name.to_string())) + }) + .collect() +} + +/// Put back the field ids a scan's output schema drops, so the reconciliation +/// can resolve its columns by id. +fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { + fn restore(field: &Field, among: &Fields) -> Field { + let Some(source) = among.iter().find(|f| f.name() == field.name()) else { + return field.clone(); + }; + let mut metadata = field.metadata().clone(); + metadata.extend(source.metadata().clone()); + let field = field.clone().with_metadata(metadata); + // A struct's children carry their own ids, and a child can be renamed + // while its parent's name does not move. + let data_type = field.data_type().clone(); + match (&data_type, source.data_type()) { + (DataType::Struct(children), DataType::Struct(source_children)) => { + let children: Vec = children + .iter() + .map(|child| restore(child, source_children)) + .collect(); + field.with_data_type(DataType::Struct(children.into())) + } + // A list's element is a field with an id of its own, and so are its + // children in turn. + (DataType::List(element), DataType::List(source_element)) => { + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::List(Arc::new(restore(element, &one)))) + } + (DataType::LargeList(element), DataType::LargeList(source_element)) => { + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::LargeList(Arc::new(restore(element, &one)))) + } + ( + DataType::FixedSizeList(element, size), + DataType::FixedSizeList(source_element, _), + ) => { + let size = *size; + let one: Fields = vec![source_element.as_ref().clone()].into(); + field.with_data_type(DataType::FixedSizeList( + Arc::new(restore(element, &one)), + size, + )) + } + _ => field, + } + } + let fields: Vec = schema + .fields() + .iter() + .map(|field| restore(field, stored.fields())) + .collect(); + Schema::new_with_metadata(fields, schema.metadata().clone()) +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 6442d43f7b0..fb442da3b58 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -3,30 +3,24 @@ //! Query planner for LSM scanner. -use std::collections::HashMap; use std::sync::Arc; -use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; -use datafusion::common::DFSchema; -use datafusion::execution::context::ExecutionProps; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, limit::GlobalLimitExec}; use datafusion::prelude::{Expr, col}; -use datafusion_physical_expr::create_physical_expr; use lance_core::Result; -use lance_core::is_system_column; use tracing::instrument; -use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; -use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; +use crate::dataset::mem_wal::TOMBSTONE; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{ - MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, ReconcileExec, + MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, }; +use super::generation::{GenerationRead, filter_above}; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, validate_projection_names, @@ -68,122 +62,6 @@ pub struct LsmScanPlanner { warmer: Option>, } -/// Apply `expr` above a source that has been reconciled, for a generation whose -/// stored names the predicate could not be run against directly. -fn filter_above(plan: Arc, expr: &Expr) -> Result> { - let schema = plan.schema(); - let df_schema = DFSchema::try_from(schema.as_ref().clone()) - .map_err(|e| lance_core::Error::internal(format!("filter schema: {e}")))?; - let props = ExecutionProps::new(); - let physical = create_physical_expr(expr, &df_schema, &props) - .map_err(|e| lance_core::Error::internal(format!("plan filter `{expr}`: {e}")))?; - Ok(Arc::new(FilterExec::try_new(physical, plan).map_err( - |e| lance_core::Error::internal(format!("filter: {e}")), - )?)) -} - -/// Whether a column of this type is rebuilt rather than taken as it stands. -/// -/// A struct's children are part of its own type, so the array is rebuilt under -/// the table's children -- which means the values a predicate would see here -/// are not the ones it will see after reconciliation. -fn is_reconstructed(data_type: &DataType) -> bool { - matches!(data_type, DataType::Struct(_)) -} - -/// What the table calls each of this generation's stored columns, keyed by the -/// stored name. -/// -/// Matched by field id, which a rename keeps. A stored column whose id the -/// table no longer declares is absent from the result: the table has dropped -/// it, and reading it would answer with data the table no longer has. -pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap { - let by_id: HashMap = table - .fields() - .iter() - .filter_map(|f| field_id_of(f).map(|id| (id, f.name().as_str()))) - .collect(); - // A caller that supplies no ids leaves only names to match on. - if by_id.is_empty() { - return stored - .fields() - .iter() - .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) - .filter(|f| table.field_with_name(f.name()).is_ok()) - .map(|f| (f.name().clone(), f.name().clone())) - .collect(); - } - stored - .fields() - .iter() - // A generation's own columns are numbered in its own schema, so their - // ids collide with whatever the table gave those numbers. They are not - // the table's columns and are never resolved to one. - .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) - .filter_map(|f| { - let id = field_id_of(f)?; - by_id - .get(&id) - .map(|name| (f.name().clone(), name.to_string())) - }) - .collect() -} - -/// `schema` with each field carrying the id `stored` gives the same name. -/// -/// A scan's output schema is built from the dataset and carries no ids, so they -/// are put back before identity is resolved against it. -fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { - fn restore(field: &Field, among: &Fields) -> Field { - let Some(source) = among.iter().find(|f| f.name() == field.name()) else { - return field.clone(); - }; - let mut metadata = field.metadata().clone(); - metadata.extend(source.metadata().clone()); - let field = field.clone().with_metadata(metadata); - // A struct's children carry their own ids, and a child can be renamed - // while its parent's name does not move. - let data_type = field.data_type().clone(); - match (&data_type, source.data_type()) { - (DataType::Struct(children), DataType::Struct(source_children)) => { - let children: Vec = children - .iter() - .map(|child| restore(child, source_children)) - .collect(); - field.with_data_type(DataType::Struct(children.into())) - } - // A list's element is a field with an id of its own, and so are its - // children in turn. - (DataType::List(element), DataType::List(source_element)) => { - let one: Fields = vec![source_element.as_ref().clone()].into(); - field.with_data_type(DataType::List(Arc::new(restore(element, &one)))) - } - (DataType::LargeList(element), DataType::LargeList(source_element)) => { - let one: Fields = vec![source_element.as_ref().clone()].into(); - field.with_data_type(DataType::LargeList(Arc::new(restore(element, &one)))) - } - ( - DataType::FixedSizeList(element, size), - DataType::FixedSizeList(source_element, _), - ) => { - let size = *size; - let one: Fields = vec![source_element.as_ref().clone()].into(); - field.with_data_type(DataType::FixedSizeList( - Arc::new(restore(element, &one)), - size, - )) - } - _ => field, - } - } - let fields: Vec = schema - .fields() - .iter() - .map(|field| restore(field, stored.fields())) - .collect(); - Schema::new_with_metadata(fields, schema.metadata().clone()) -} - impl LsmScanPlanner { /// Create a new planner. pub fn new( @@ -446,34 +324,6 @@ impl LsmScanPlanner { Arc::new(Schema::new(fields)) } - /// Build scan plan for a single data source. - /// What one generation's scan should produce once reconciled: the columns - /// this query needs, named as the table names them, plus the ones the - /// generation carries of its own. - /// - /// The same for every generation, and taken from the table rather than from - /// whichever source the collector ordered first. - fn generation_target( - &self, - wanted: &[String], - source: &Schema, - names: &HashMap, - ) -> SchemaRef { - let mut fields: Vec = wanted - .iter() - .filter_map(|name| self.identity_schema.field_with_name(name).ok().cloned()) - .collect(); - // A generation's own columns are not the table's, so they pass through - // as the generation has them. - for field in source.fields() { - let is_the_tables = names.contains_key(field.name()); - if !is_the_tables && fields.iter().all(|f| f.name() != field.name()) { - fields.push(field.as_ref().clone()); - } - } - Arc::new(Schema::new(fields)) - } - async fn build_source_scan( &self, source: &LsmDataSource, @@ -517,58 +367,34 @@ impl LsmScanPlanner { .await?; let mut scanner = dataset.scan(); - // What the table calls each of this generation's columns, by - // field id: a rename changes the name and keeps the id. - let stored = arrow_schema_with_field_ids(dataset.schema()); - let names = stored_names(&stored, &self.identity_schema); - // Asked of this generation under its own names, so an older // file is only asked for columns it has. A column it never had // is filled in after the scan. - let mut wanted = + let wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); + let mut generation = GenerationRead::new( + dataset.schema(), + Arc::clone(&self.identity_schema), + self.pk_columns.clone(), + wanted, + ); // A predicate this generation cannot answer as written runs // after reconciliation, reading its columns from this scan -- // so they have to be in it whether the caller asked or not. - // A predicate can run here only against columns this generation - // stores under the table's own name and shape. A nested column - // names its parent in the reference -- `info.a` refers to - // `info` -- and the parent's name does not move when a child is - // renamed, so a name check alone would push the predicate down - // to be evaluated against the child names this generation has. - // Anything reconstructed is answered after reconciliation. - let answerable = filter.is_none_or(|expr| { - expr.column_refs().iter().all(|c| { - names.get(c.name.as_str()) == Some(&c.name) - && stored - .field_with_name(&c.name) - .is_ok_and(|f| !is_reconstructed(f.data_type())) - }) - }); + let answerable = filter.is_none_or(|expr| generation.can_answer(expr)); if let Some(expr) = filter.filter(|_| !answerable) { for column in expr.column_refs() { - if !wanted.contains(&column.name) { - wanted.push(column.name.clone()); - } + generation.also_produce(&column.name); } } - let cols: Vec<&str> = wanted - .iter() - .filter_map(|name| { - names - .iter() - .find(|(_, table_name)| *table_name == name) - .map(|(stored_name, _)| stored_name.as_str()) - }) - .collect(); - scanner.project(&cols)?; + scanner.project(&generation.stored_projection())?; scanner.with_row_address(); // Drop tombstones: fold `NOT _tombstone` into the predicate so // it runs before the pushdown limit (counting only live rows). // The older real row a tombstone supersedes is dropped by the // cross-gen block-list, not by this filter. Gen written before - // deletes existed lack the column → no fold, nothing to drop. + // deletes existed lack the column -> no fold, nothing to drop. let caller_filter = if answerable { filter } else { None }; let folded; let effective: Option<&Expr> = if dataset.schema().field(TOMBSTONE).is_some() { @@ -586,15 +412,7 @@ impl LsmScanPlanner { scanner.limit(Some(fetch as i64), None)?; } - let scan = scanner.create_plan().await?; - // Planned against what the scan actually produces, which is the - // projection plus `_rowaddr`, and carrying the ids the dataset - // gives those columns. - let source = with_ids_from(&scan.schema(), &stored); - let target = self.generation_target(&wanted, &source, &names); - let plan = Plan::resolve(&source, &target, &self.pk_columns)?; - let reconciled: Arc = - Arc::new(ReconcileExec::new(scan, Arc::new(plan))); + let reconciled = generation.reconcile(scanner.create_plan().await?)?; match filter { Some(expr) if !answerable => filter_above(reconciled, expr), _ => Ok(reconciled), diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index 9aea4e9f4e7..d5e2656a054 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -38,6 +38,8 @@ use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, force_schema, null_columns, project_to_canonical, validate_projection_names, wants_row_address, wants_row_id, }; +use super::generation::GenerationRead; +use crate::dataset::mem_wal::reconcile::relabel_to; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; use crate::session::Session; use lance_io::object_store::ObjectStoreParams; @@ -108,6 +110,10 @@ pub struct LsmPointLookupPlanner { /// Prefix of the in-memory memtables this planner reads. Applies to the fast /// BTree probe and the plan fallback alike, so both resolve a key the same. visibility: MemTableVisibility, + /// `base_schema` with each field's id, which is what resolves a sealed + /// generation's columns to the table's. Defaults to `base_schema`, which + /// carries them when the caller built it from a Lance schema. + identity_schema: SchemaRef, } impl LsmPointLookupPlanner { @@ -127,7 +133,7 @@ impl LsmPointLookupPlanner { Self { collector, pk_columns, - base_schema, + base_schema: Arc::clone(&base_schema), bloom_filters: std::collections::HashMap::new(), session: None, store_params: None, @@ -136,9 +142,16 @@ impl LsmPointLookupPlanner { none_target, task_ctx: SessionContext::new().task_ctx(), visibility: MemTableVisibility::Published, + identity_schema: base_schema, } } + /// Supply the table's field ids, when `base_schema` was built without them. + pub fn with_identity_schema(mut self, identity_schema: SchemaRef) -> Self { + self.identity_schema = identity_schema; + self + } + /// Read the in-memory memtables at `visibility`. See /// [`MemTableVisibility::Indexed`] for when a wider bound is sound. pub fn with_visibility(mut self, visibility: MemTableVisibility) -> Self { @@ -683,13 +696,31 @@ impl LsmPointLookupPlanner { ) .await?; let mut scanner = dataset.scan(); - // Carry `_tombstone` through so the post-coalesce filter can drop - // a deleted key (gen written before deletes existed lack it → - // `project_to_carry` synthesizes `false`). + // A sealed generation holds the names the table had when it was + // sealed, so the projection, the key filter and the output all + // go through the same resolution the scanner uses. let cols = cols_with_tombstone(&cols, dataset.schema().field(TOMBSTONE).is_some()); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; - scanner.filter_expr(filter.clone()); - Box::pin(scanner.create_plan()).await? + let generation = GenerationRead::new( + dataset.schema(), + Arc::clone(&self.identity_schema), + self.pk_columns.clone(), + cols, + ); + // Every generation stores every primary key column -- a key + // cannot be added or dropped -- so the key filter always moves. + let stored_filter = generation.to_stored(filter).ok_or_else(|| { + lance_core::Error::internal(format!( + "point lookup: `{filter}` names a column generation {} does not store", + source.generation() + )) + })?; + scanner.project(&generation.stored_projection())?; + scanner.filter_expr(stored_filter); + let scan = Box::pin(scanner.create_plan()).await?; + // `_tombstone` is carried through so the post-coalesce filter + // can drop a deleted key; a generation written before deletes + // existed lacks it and `project_to_carry` synthesizes `false`. + generation.reconcile(scan)? } LsmDataSource::ActiveMemTable { batch_store, @@ -1007,11 +1038,13 @@ fn gather_rows( // Single row: zero-copy `slice` (the common point-lookup case, and // measurably faster than `take` — copying regressed single-thread // ~30% with no N-thread gain). Multiple rows: one vectorized `take`. - match &indices { - None => Ok(col.slice(rows[0] as usize, 1)), - Some(idxs) => arrow_select::take::take(col.as_ref(), idxs, None) - .map_err(lance_core::Error::from), - } + let picked = match &indices { + None => col.slice(rows[0] as usize, 1), + Some(idxs) => arrow_select::take::take(col.as_ref(), idxs, None)?, + }; + // The memtable stores each field's id inside a nested column's own + // type; the caller asked for the table's plain one. + relabel_to(&picked, f.data_type()) }) .collect::>>()?; Ok(RecordBatch::try_new(target.clone(), cols)?) diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index d258f299d7c..e5a0c99d261 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -28,13 +28,12 @@ use crate::io::exec::TakeExec; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; -use super::planner::stored_names; +use super::generation::{GenerationRead, filter_above}; use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, validate_projection_names, wants_row_id, }; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; -use crate::dataset::mem_wal::arrow_schema_with_field_ids; use crate::session::Session; use lance_io::object_store::ObjectStoreParams; @@ -132,15 +131,14 @@ impl LsmVectorSearchPlanner { collector: LsmDataSourceCollector, pk_columns: Vec, base_schema: SchemaRef, - identity_schema: SchemaRef, vector_column: String, distance_type: lance_linalg::distance::DistanceType, ) -> Self { Self { collector, pk_columns, + identity_schema: base_schema.clone(), base_schema, - identity_schema, vector_column, distance_type, dataset: None, @@ -182,6 +180,16 @@ impl LsmVectorSearchPlanner { } /// Set the session used to open SSTables. + /// The table's schema carrying each field's id, which is what resolves a + /// generation's stored columns to the table's. + /// + /// Defaults to the base schema, so a caller that has no ids to give is + /// matched by name as it was. + pub fn with_identity_schema(mut self, schema: SchemaRef) -> Self { + self.identity_schema = schema; + self + } + pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); self @@ -506,29 +514,43 @@ impl LsmVectorSearchPlanner { // the table's name while the file still holds the old one, so // projecting the table's names would ask for a column that is // not there. - let stored = arrow_schema_with_field_ids(dataset.schema()); - let names = stored_names(&stored, &self.identity_schema); let wanted = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - let cols: Vec<&str> = wanted - .iter() - .filter_map(|name| { - names - .iter() - .find(|(_, table_name)| *table_name == name) - .map(|(stored_name, _)| stored_name.as_str()) - }) - .collect(); - scanner.project(&cols)?; - if let Some(ref filter) = self.filter { + let mut generation = GenerationRead::new( + dataset.schema(), + Arc::clone(&self.identity_schema), + self.pk_columns.clone(), + wanted, + ); + // The index is on this generation's own column, under the name + // it had when the generation was sealed. + let Some(vector_column) = generation.stored_name(&self.vector_column) else { + // The table dropped the column the search names, so this + // generation has no candidates to offer. + return self.empty_plan(projection); + }; + let vector_column = vector_column.to_string(); + // A predicate this generation cannot answer as written runs + // above the reconciliation, where the columns it names exist. + // The search's own top-k has already run by then, so the arm + // can come back short -- which is right: those rows have no + // value for a column sealed before it existed. + let pushed = self.filter.as_ref().map(|f| (f, generation.to_stored(f))); + if let Some((expr, None)) = pushed { + for column in expr.column_refs() { + generation.also_produce(&column.name); + } + } + scanner.project(&generation.stored_projection())?; + if let Some((_, Some(ref stored))) = pushed { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the top-k. - scanner.filter_expr(filter.clone()); + scanner.filter_expr(stored.clone()); scanner.prefilter(true); } // No `with_row_id/address`: per-source IDs would collide with base. let query_arr = single_query_array(query_vector); - scanner.nearest(&self.vector_column, query_arr.as_ref(), k)?; + scanner.nearest(&vector_column, query_arr.as_ref(), k)?; scanner.distance_range(self.distance_range.0, self.distance_range.1); scanner.nprobes(nprobes); scanner.distance_metric(self.distance_type); @@ -536,7 +558,11 @@ impl LsmVectorSearchPlanner { scanner.ef(ef); } scanner.fast_search(); - scanner.create_plan().await + let reconciled = generation.reconcile(scanner.create_plan().await?)?; + match pushed { + Some((expr, None)) => filter_above(reconciled, expr), + _ => Ok(reconciled), + } } LsmDataSource::ActiveMemTable { batch_store, diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 4a9b8e61a40..1456fc9cd1a 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -4581,8 +4581,10 @@ mod tests { use super::*; use crate::dataset::mem_wal::test_util::failing_memory_store; use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, Int64Array, StringArray}; + use arrow_schema::Field as ArrowField; use arrow_schema::{DataType, Field}; use lance_core::FenceReason; + use lance_core::datatypes::LANCE_FIELD_ID_KEY; use rstest::rstest; use std::sync::atomic::AtomicUsize; use tempfile::TempDir; @@ -4812,11 +4814,10 @@ mod tests { "the error should name the missing key: {error}" ); } - - /// A widened type is cast, matching what `alter_columns` did to the rows - /// already in the base table. + /// A column whose type has moved is refused rather than cast: a table with a + /// MemWAL does not accept a retype, so this is a disagreement to surface. #[test] - fn test_conform_casts_a_widened_column() { + fn test_conform_refuses_a_column_whose_type_moved() { let widened = ArrowSchema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), @@ -4838,15 +4839,12 @@ mod tests { ) .unwrap(); - let out = conform_to_storage_schema(narrow, &storage, &["id".to_string()]).unwrap(); - assert_eq!(out.schema(), storage); - let counts = out - .column_by_name("count") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(counts.value(0), 7, "the value survives the widening"); + let err = conform_to_storage_schema(narrow, &storage, &["id".to_string()]) + .expect_err("a column's type cannot change"); + assert!( + err.to_string().contains("count"), + "the refusal should name the column, got: {err}" + ); } /// A cast that would lose the value is an error, not a column of nulls. From 83d76895653fd667137ff620e49ed58eb87fea15 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 21:19:03 -0400 Subject: [PATCH 25/48] fix(mem_wal): resolve the full-text index's column per generation too The searched column's name is settled once per query, from the table's schema, and then used to find each source's index. A generation sealed before a rename holds the index under the old name, so the granularity lookup failed the whole query where the arm that reads the rows had already been taught to translate. One sealed before the column existed has no index on it and now offers no granularity rather than erroring. GenerationRead gains unit tests for the rules the read paths depend on: which name a column is asked for, when a predicate can be pushed down and under which name, and that a generation's own columns never answer for one of the table's. --- .../src/dataset/mem_wal/scanner/fts_search.rs | 30 ++- .../src/dataset/mem_wal/scanner/generation.rs | 192 ++++++++++++++++++ 2 files changed, 215 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index f995c2c2757..dd94b1c7f26 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -693,14 +693,30 @@ impl LsmFtsSearchPlanner { self.warmer.as_ref(), ) .await?; - if index_params.is_empty() { - index_params = indexed_fts_index_params(&dataset, column).await?; + // The index is on this generation's own column, under the + // name it had when the generation was sealed. A generation + // sealed before the column existed has no index on it and + // offers no granularity. + let generation = GenerationRead::new( + dataset.schema(), + Arc::clone(&self.identity_schema), + self.pk_columns.clone(), + Vec::new(), + ); + match generation.stored_name(column) { + None => Vec::new(), + Some(stored_column) => { + if index_params.is_empty() { + index_params = + indexed_fts_index_params(&dataset, stored_column).await?; + } + indexed_fts_document_granularities(&dataset, stored_column) + .await? + .into_iter() + .map(|(_, document_granularity)| document_granularity) + .collect::>() + } } - indexed_fts_document_granularities(&dataset, column) - .await? - .into_iter() - .map(|(_, document_granularity)| document_granularity) - .collect::>() } LsmDataSource::ActiveMemTable { index_store, .. } => { index_store.fts_document_granularities_by_column(column) diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index fd4a9940be7..800030d6c3f 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -327,3 +327,195 @@ fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { .collect(); Schema::new_with_metadata(fields, schema.metadata().clone()) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow_schema::Fields; + use datafusion::prelude::{col, lit}; + use lance_core::datatypes::Schema as LanceSchema; + + /// An Arrow field carrying a Lance field id, as a generation's schema and + /// the table's both do. + fn with_id(name: &str, data_type: DataType, id: i32) -> Field { + Field::new(name, data_type, true).with_metadata( + [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] + .into_iter() + .collect(), + ) + } + + fn schema(fields: Vec) -> SchemaRef { + Arc::new(Schema::new(fields)) + } + + /// `GenerationRead` resolves against a generation's *Lance* schema, which + /// is where the stored ids come from. + fn generation(stored: SchemaRef, table: SchemaRef, wanted: &[&str]) -> GenerationRead { + let lance = LanceSchema::try_from(stored.as_ref()).expect("a lance schema"); + GenerationRead::new( + &lance, + table, + vec!["id".to_string()], + wanted.iter().map(|s| s.to_string()).collect(), + ) + } + + /// The generation was sealed as `value`; the table has since renamed it. + fn renamed() -> GenerationRead { + generation( + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("value", DataType::Int64, 1), + ]), + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("amount", DataType::Int64, 1), + ]), + &["id", "amount"], + ) + } + + #[test] + fn a_renamed_column_is_asked_for_under_the_name_the_generation_has() { + assert_eq!(renamed().stored_projection(), vec!["id", "value"]); + assert_eq!(renamed().stored_name("amount"), Some("value")); + } + + #[test] + fn a_column_the_generation_never_stored_is_left_out_of_the_projection() { + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0)]), + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("added", DataType::Int64, 7), + ]), + &["id", "added"], + ); + assert_eq!(read.stored_projection(), vec!["id"]); + assert_eq!(read.stored_name("added"), None); + } + + /// A generation numbers its own columns in its own schema, so `_tombstone` + /// carries an id that collides with whatever the table gave that number. + #[test] + fn a_system_column_never_answers_for_one_of_the_tables() { + let read = generation( + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id(TOMBSTONE, DataType::Boolean, 7), + ]), + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("added", DataType::Int64, 7), + ]), + &["id", "added", TOMBSTONE], + ); + assert_eq!(read.stored_name("added"), None, "not the tombstone's id"); + assert_eq!( + read.stored_projection(), + vec!["id", TOMBSTONE], + "the tombstone is still asked for, under its own name" + ); + } + + #[test] + fn a_predicate_naming_a_renamed_column_is_rewritten_to_the_stored_name() { + let read = renamed(); + assert!( + !read.can_answer(&col("amount").eq(lit(1i64))), + "the generation has no `amount`" + ); + assert_eq!( + read.to_stored(&col("amount").eq(lit(1i64))), + Some(col("value").eq(lit(1i64))), + ); + } + + #[test] + fn a_predicate_the_generation_can_answer_as_written_is_left_alone() { + let read = renamed(); + let expr = col("id").eq(lit(1i64)); + assert!(read.can_answer(&expr)); + assert_eq!(read.to_stored(&expr), Some(expr)); + } + + /// A predicate on a column the generation never stored cannot be pushed + /// down; it belongs above the reconciliation, where the column exists as + /// nulls. + #[test] + fn a_predicate_naming_a_column_the_generation_lacks_is_not_pushable() { + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0)]), + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("added", DataType::Int64, 7), + ]), + &["id", "added"], + ); + let expr = col("added").eq(lit(1i64)); + assert!(!read.can_answer(&expr)); + assert_eq!(read.to_stored(&expr), None); + } + + /// A nested reference names the parent, and a parent's name does not move + /// when a child is renamed -- so pushing it down would evaluate it against + /// child names the table does not have. + #[test] + fn a_predicate_on_a_nested_column_is_never_pushed_down() { + let nested = |child: &str| { + with_id( + "info", + DataType::Struct(Fields::from(vec![with_id(child, DataType::Int64, 2)])), + 1, + ) + }; + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0), nested("c")]), + schema(vec![with_id("id", DataType::Int64, 0), nested("d")]), + &["id", "info"], + ); + let expr = col("info").is_not_null(); + assert!(!read.can_answer(&expr)); + assert_eq!(read.to_stored(&expr), None); + } + + /// A list of structs is rebuilt too, so its stored shape is not the shape a + /// pushed-down predicate would expect. + #[test] + fn a_list_is_reconstructed_only_when_its_element_is() { + let struct_element = DataType::Struct(Fields::from(vec![Field::new( + "c", + DataType::Int64, + true, + )])); + assert!(is_reconstructed(&DataType::List(Arc::new(Field::new( + "item", + struct_element, + true + ))))); + assert!(!is_reconstructed(&DataType::List(Arc::new(Field::new( + "item", + DataType::Int64, + true + ))))); + } + + /// With no ids to match on, the table's own names are the only link -- the + /// behaviour a caller that supplies no identity schema gets. + #[test] + fn a_table_without_ids_matches_by_name() { + let stored = Schema::new(vec![ + with_id("id", DataType::Int64, 0), + with_id("value", DataType::Int64, 1), + with_id(TOMBSTONE, DataType::Boolean, 2), + ]); + let table = Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Int64, true), + ]); + let names = stored_names(&stored, &table); + assert_eq!(names.get("value"), Some(&"value".to_string())); + assert_eq!(names.get(TOMBSTONE), None, "not one of the table's columns"); + } +} From 64b1b45859e264fcac8c96e27284272dc2617640 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 21:35:56 -0400 Subject: [PATCH 26/48] refactor(mem_wal): push a predicate down under the names a generation has The scan refused to push a predicate naming a renamed column and ran it above the reconciliation instead, while the two search paths rewrote it and pushed it down. Rewriting is the better half of that split -- the generation's own scan can answer it -- and it leaves one rule instead of two, so `can_answer` goes and every arm asks `to_stored` the same way. A predicate still runs above the reconciliation when it cannot be moved at all: a column the generation was sealed before, or a nested one. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 10 +-- .../mem_wal/scanner/exec/schema_relabel.rs | 2 +- .../src/dataset/mem_wal/scanner/fts_search.rs | 18 +++-- .../src/dataset/mem_wal/scanner/generation.rs | 70 +++++++------------ .../src/dataset/mem_wal/scanner/planner.rs | 38 +++++----- .../dataset/mem_wal/scanner/point_lookup.rs | 8 +-- .../dataset/mem_wal/scanner/vector_search.rs | 18 +++-- 7 files changed, 77 insertions(+), 87 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 41f60e4ce8e..b2155b4453a 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -44,7 +44,7 @@ pub(crate) fn field_id_of(field: &ArrowField) -> Option { /// /// Ids belong to the stored schema, where identity has to survive a rename. A /// caller's batch carries none, and Arrow compares a struct's children by their -/// full field -- metadata included -- so a stamped schema would reject it. +/// full field — metadata included — so a stamped schema would reject it. /// One stored column under the type the caller declared. /// /// A field id lives inside a nested column's own Arrow type, so the memtable's @@ -132,8 +132,8 @@ pub struct Plan { impl Plan { /// The same plan emitting the table's plain Arrow schema. /// - /// Resolution needs field ids on the target -- a rename keeps the id and - /// moves the name -- but a reader is handed the table's schema, which does + /// Resolution needs field ids on the target — a rename keeps the id and + /// moves the name — but a reader is handed the table's schema, which does /// not carry them. They live inside a nested column's own type, so a struct /// built to the id-carrying target is a different Arrow type from the one /// the caller declared. Only replay, which writes back into the memtable's @@ -240,7 +240,7 @@ fn resolve_field( }); // Identity is the field id where both sides carry one. A name is not: a // rename moves the name and leaves the id, so a source column of the same - // name under a *different* id is a different column -- one dropped and + // name under a *different* id is a different column — one dropped and // another added under its name, whose values the table no longer has. // // A name match is right only where identity is absent: a batch a caller has @@ -280,7 +280,7 @@ fn resolve_field( let source = &source_fields[index]; // A nested column's children are part of the array's own type, metadata // included, so the array is rebuilt under the target's children even when - // nothing about them moved -- otherwise the batch disagrees with the schema + // nothing about them moved — otherwise the batch disagrees with the schema // it is built under. The leaves are reused, so it costs a pointer copy. if is_nested(field.data_type()) { return Ok(Source::Nested( diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index ddf1d017d73..e629bc3e293 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -8,6 +8,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use crate::dataset::mem_wal::reconcile::relabel_to; use arrow_array::{RecordBatch, RecordBatchOptions}; use arrow_schema::SchemaRef; use datafusion::error::{DataFusionError, Result as DFResult}; @@ -18,7 +19,6 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use futures::{Stream, StreamExt}; -use crate::dataset::mem_wal::reconcile::relabel_to; /// Re-labels every batch to an exact target schema, leaving the arrays /// untouched. `ProjectionExec` cannot: DataFusion derives output nullability diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index dd94b1c7f26..e0a5298120e 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -961,16 +961,20 @@ impl LsmFtsSearchPlanner { // A predicate this generation cannot answer as written runs // above the reconciliation, where the columns it names exist. // The BM25 top-k has already run by then, so the arm can come - // back short -- which is right: those rows have no value for a + // back short — which is right: those rows have no value for a // column sealed before it existed. - let pushed = self.filter.as_ref().map(|f| (f, generation.to_stored(f))); - if let Some((expr, None)) = pushed { + let stored_filter = self + .filter + .as_ref() + .and_then(|expr| generation.to_stored(expr)); + let above = self.filter.as_ref().filter(|_| stored_filter.is_none()); + if let Some(expr) = above { for column in expr.column_refs() { generation.also_produce(&column.name); } } scanner.project(&generation.stored_projection())?; - if let Some((_, Some(ref stored))) = pushed { + if let Some(ref stored) = stored_filter { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the BM25 top-k. scanner.filter_expr(stored.clone()); @@ -984,9 +988,9 @@ impl LsmFtsSearchPlanner { } scanner.full_text_search(bound_query)?; let reconciled = generation.reconcile(scanner.create_plan().await?)?; - match pushed { - Some((expr, None)) => filter_above(reconciled, expr), - _ => Ok(reconciled), + match above { + Some(expr) => filter_above(reconciled, expr), + None => Ok(reconciled), } } LsmDataSource::ActiveMemTable { diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index 800030d6c3f..5270e8f31e1 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -21,12 +21,12 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::filter::FilterExec; use datafusion::prelude::Expr; use datafusion_physical_expr::create_physical_expr; +use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::is_system_column; use lance_core::{Error, Result}; use super::exec::ReconcileExec; use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; -use lance_core::datatypes::LANCE_FIELD_ID_KEY; use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// One sealed generation, read under the table's schema. @@ -34,7 +34,7 @@ use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// Built from the generation's own schema and the table's id-carrying schema. /// It answers three questions, in the order a scan needs them: what to project /// from the file ([`Self::stored_projection`]), whether a predicate can be -/// pushed into it ([`Self::can_answer`] / [`Self::to_stored`]), and how to +/// pushed into it and under which names ([`Self::to_stored`]), and how to /// bring the result back to the table's names and shapes /// ([`Self::reconcile`]). pub(super) struct GenerationRead { @@ -92,7 +92,10 @@ impl GenerationRead { pub(super) fn stored_projection(&self) -> Vec<&str> { self.wanted .iter() - .filter_map(|name| self.stored_name(name).or_else(|| self.stored_system_column(name))) + .filter_map(|name| { + self.stored_name(name) + .or_else(|| self.stored_system_column(name)) + }) .collect() } @@ -106,26 +109,15 @@ impl GenerationRead { .map(|f| f.name().as_str()) } - /// Whether `expr` can be pushed into this generation's scan. - /// - /// It can when every column it names is stored under the table's own name - /// and shape. A nested column fails the shape half: a reference names the - /// parent (`info.a` refers to `info`) and a parent's name does not move - /// when a child is renamed, so the pushed-down predicate would be evaluated - /// against child names this generation has and the table does not. - pub(super) fn can_answer(&self, expr: &Expr) -> bool { - expr.column_refs().iter().all(|c| { - self.names.get(c.name.as_str()) == Some(&c.name) - && self - .stored - .field_with_name(&c.name) - .is_ok_and(|f| !is_reconstructed(f.data_type())) - }) - } - /// `expr` with each column reference moved to the name this generation - /// stores it under. `None` when a column it names is not stored here at - /// all, or is reconstructed — then the predicate belongs above the scan. + /// stores it under, so it can be pushed into the generation's own scan. + /// + /// `None` when it cannot be: a column the generation does not store, or a + /// nested one. A nested reference names the parent (`info.a` refers to + /// `info`) and a parent's name does not move when a child is renamed, so a + /// pushed-down predicate would be evaluated against child names this + /// generation has and the table does not. Either way the predicate belongs + /// above the reconciliation, where the columns it names exist. pub(super) fn to_stored(&self, expr: &Expr) -> Option { let pushable = expr.column_refs().iter().all(|c| { self.stored_name(&c.name).is_some_and(|stored| { @@ -171,8 +163,8 @@ impl GenerationRead { /// `source` with the field ids of everything that is not one of the /// table's columns removed. /// - /// A generation numbers its own columns in its own schema -- `_tombstone`, - /// and anything it holds that the table has since dropped -- so those ids + /// A generation numbers its own columns in its own schema — `_tombstone`, + /// and anything it holds that the table has since dropped — so those ids /// collide with whatever the table gave those numbers. Left in place, a /// column added to the table resolves to whichever of them happens to share /// its id. @@ -212,8 +204,8 @@ impl GenerationRead { } } -/// Run `expr` above `plan`, for a predicate the generation could not answer -/// as written. +/// Run `expr` above `plan`, for a predicate that could not be pushed into the +/// generation's own scan. pub(super) fn filter_above( plan: Arc, expr: &Expr, @@ -422,10 +414,6 @@ mod tests { #[test] fn a_predicate_naming_a_renamed_column_is_rewritten_to_the_stored_name() { let read = renamed(); - assert!( - !read.can_answer(&col("amount").eq(lit(1i64))), - "the generation has no `amount`" - ); assert_eq!( read.to_stored(&col("amount").eq(lit(1i64))), Some(col("value").eq(lit(1i64))), @@ -433,10 +421,9 @@ mod tests { } #[test] - fn a_predicate_the_generation_can_answer_as_written_is_left_alone() { + fn a_predicate_naming_a_column_that_did_not_move_is_left_alone() { let read = renamed(); let expr = col("id").eq(lit(1i64)); - assert!(read.can_answer(&expr)); assert_eq!(read.to_stored(&expr), Some(expr)); } @@ -453,13 +440,11 @@ mod tests { ]), &["id", "added"], ); - let expr = col("added").eq(lit(1i64)); - assert!(!read.can_answer(&expr)); - assert_eq!(read.to_stored(&expr), None); + assert_eq!(read.to_stored(&col("added").eq(lit(1i64))), None); } /// A nested reference names the parent, and a parent's name does not move - /// when a child is renamed -- so pushing it down would evaluate it against + /// when a child is renamed — so pushing it down would evaluate it against /// child names the table does not have. #[test] fn a_predicate_on_a_nested_column_is_never_pushed_down() { @@ -475,20 +460,15 @@ mod tests { schema(vec![with_id("id", DataType::Int64, 0), nested("d")]), &["id", "info"], ); - let expr = col("info").is_not_null(); - assert!(!read.can_answer(&expr)); - assert_eq!(read.to_stored(&expr), None); + assert_eq!(read.to_stored(&col("info").is_not_null()), None); } /// A list of structs is rebuilt too, so its stored shape is not the shape a /// pushed-down predicate would expect. #[test] fn a_list_is_reconstructed_only_when_its_element_is() { - let struct_element = DataType::Struct(Fields::from(vec![Field::new( - "c", - DataType::Int64, - true, - )])); + let struct_element = + DataType::Struct(Fields::from(vec![Field::new("c", DataType::Int64, true)])); assert!(is_reconstructed(&DataType::List(Arc::new(Field::new( "item", struct_element, @@ -501,7 +481,7 @@ mod tests { ))))); } - /// With no ids to match on, the table's own names are the only link -- the + /// With no ids to match on, the table's own names are the only link — the /// behaviour a caller that supplies no identity schema gets. #[test] fn a_table_without_ids_matches_by_name() { diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index fb442da3b58..fdaf8f7348b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -17,9 +17,7 @@ use crate::dataset::mem_wal::TOMBSTONE; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; -use super::exec::{ - MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN, -}; +use super::exec::{MEMTABLE_GEN_COLUMN, MemtableGenTagExec, PkBlockFilterExec, ROW_ADDRESS_COLUMN}; use super::generation::{GenerationRead, filter_above}; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, @@ -245,7 +243,7 @@ impl LsmScanPlanner { // column was added does not carry it. `UnionExec` requires schema // equality and does not reconcile. // - // The base arm is the authority when it is here -- it is the only source + // The base arm is the authority when it is here — it is the only source // the schema change was applied to. Otherwise the newest generation is, // and sources arrive generation-DESC, so it is the first of them. let target = source_plans @@ -378,11 +376,15 @@ impl LsmScanPlanner { self.pk_columns.clone(), wanted, ); - // A predicate this generation cannot answer as written runs - // after reconciliation, reading its columns from this scan -- - // so they have to be in it whether the caller asked or not. - let answerable = filter.is_none_or(|expr| generation.can_answer(expr)); - if let Some(expr) = filter.filter(|_| !answerable) { + // A predicate the generation can answer is pushed into its scan + // under the names it has. One it cannot — because it names a + // column sealed before it existed, or a nested one — runs above + // the reconciliation instead, reading its columns from this + // scan, so they have to be in it whether the caller asked or + // not. + let stored_filter = filter.and_then(|expr| generation.to_stored(expr)); + let above = filter.filter(|_| stored_filter.is_none()); + if let Some(expr) = above { for column in expr.column_refs() { generation.also_produce(&column.name); } @@ -393,29 +395,29 @@ impl LsmScanPlanner { // Drop tombstones: fold `NOT _tombstone` into the predicate so // it runs before the pushdown limit (counting only live rows). // The older real row a tombstone supersedes is dropped by the - // cross-gen block-list, not by this filter. Gen written before - // deletes existed lack the column -> no fold, nothing to drop. - let caller_filter = if answerable { filter } else { None }; + // cross-gen block-list, not by this filter. A generation written + // before deletes existed lacks the column, so nothing is folded + // and there is nothing to drop. let folded; let effective: Option<&Expr> = if dataset.schema().field(TOMBSTONE).is_some() { - folded = fold_not_tombstone(caller_filter); + folded = fold_not_tombstone(stored_filter.as_ref()); Some(&folded) } else { - caller_filter + stored_filter.as_ref() }; if let Some(expr) = effective { scanner.filter_expr(expr.clone()); } // A limit under a filter that has not run would cut rows the // filter never saw. - if let Some(fetch) = fetch.filter(|_| answerable) { + if let Some(fetch) = fetch.filter(|_| above.is_none()) { scanner.limit(Some(fetch as i64), None)?; } let reconciled = generation.reconcile(scanner.create_plan().await?)?; - match filter { - Some(expr) if !answerable => filter_above(reconciled, expr), - _ => Ok(reconciled), + match above { + Some(expr) => filter_above(reconciled, expr), + None => Ok(reconciled), } } LsmDataSource::ActiveMemTable { diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index d5e2656a054..0fcc1a59ebb 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -34,13 +34,13 @@ use crate::dataset::mem_wal::{TOMBSTONE, relax_non_pk_nullability}; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{BloomFilterGuardExec, CoalesceFirstExec, compute_pk_hash_from_scalars}; +use super::generation::GenerationRead; use super::projection::{ DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, force_schema, null_columns, project_to_canonical, validate_projection_names, wants_row_address, wants_row_id, }; -use super::generation::GenerationRead; -use crate::dataset::mem_wal::reconcile::relabel_to; use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; +use crate::dataset::mem_wal::reconcile::relabel_to; use crate::session::Session; use lance_io::object_store::ObjectStoreParams; @@ -706,8 +706,8 @@ impl LsmPointLookupPlanner { self.pk_columns.clone(), cols, ); - // Every generation stores every primary key column -- a key - // cannot be added or dropped -- so the key filter always moves. + // Every generation stores every primary key column — a key + // cannot be added or dropped — so the key filter always moves. let stored_filter = generation.to_stored(filter).ok_or_else(|| { lance_core::Error::internal(format!( "point lookup: `{filter}` names a column generation {} does not store", diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index e5a0c99d261..316829cb455 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -533,16 +533,20 @@ impl LsmVectorSearchPlanner { // A predicate this generation cannot answer as written runs // above the reconciliation, where the columns it names exist. // The search's own top-k has already run by then, so the arm - // can come back short -- which is right: those rows have no + // can come back short — which is right: those rows have no // value for a column sealed before it existed. - let pushed = self.filter.as_ref().map(|f| (f, generation.to_stored(f))); - if let Some((expr, None)) = pushed { + let stored_filter = self + .filter + .as_ref() + .and_then(|expr| generation.to_stored(expr)); + let above = self.filter.as_ref().filter(|_| stored_filter.is_none()); + if let Some(expr) = above { for column in expr.column_refs() { generation.also_produce(&column.name); } } scanner.project(&generation.stored_projection())?; - if let Some((_, Some(ref stored))) = pushed { + if let Some(ref stored) = stored_filter { // See the base arm: `prefilter(true)` makes this a true // prefilter rather than a lossy post-filter on the top-k. scanner.filter_expr(stored.clone()); @@ -559,9 +563,9 @@ impl LsmVectorSearchPlanner { } scanner.fast_search(); let reconciled = generation.reconcile(scanner.create_plan().await?)?; - match pushed { - Some((expr, None)) => filter_above(reconciled, expr), - _ => Ok(reconciled), + match above { + Some(expr) => filter_above(reconciled, expr), + None => Ok(reconciled), } } LsmDataSource::ActiveMemTable { From 554b2b498b9f7e78fa82e52ef294ec6a157fb203 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 22:30:35 -0400 Subject: [PATCH 27/48] fix(mem_wal): make the reconciliation's contracts hold off the happy path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three rules the sealed-generation reader had only half right, each of which fires without any schema change at all: - The relabel replaced an array's outer type and left its children alone, so a `Struct>` was rejected: Arrow validates a container against the child types its own type declares, and a field id lives inside a nested type at every level. It now recurses. - Reconciliation produced the table's declared nullability. A tombstone is null in everything but the key, and the point lookup carries tombstones through on purpose, so a deleted key with a required payload failed validation before anything could drop it. The intermediate schema now takes nullability from the source — which already widens non-key columns for exactly this reason — and the public boundary restores the table's, after the tombstones are gone. - A predicate on any nested column was held above the reconciliation, on the assumption that a nested name may have moved. Whether it moved is knowable: the shapes are compared, so an ordinary predicate on an ordinary struct is pushed into the generation's own scan, where the search can apply it before it truncates to top-k rather than after. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 92 +++++++++++++++++-- .../src/dataset/mem_wal/scanner/generation.rs | 87 ++++++++++-------- 2 files changed, 134 insertions(+), 45 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index b2155b4453a..df88fb18243 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -25,6 +25,7 @@ use arrow_array::{ Array, ArrayRef, BooleanArray, FixedSizeListArray, GenericListArray, RecordBatch, RecordBatchOptions, StructArray, }; +use arrow::array::ArrayData; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::{Error, Result}; @@ -55,23 +56,39 @@ pub(crate) fn relabel_to(column: &ArrayRef, data_type: &DataType) -> Result Result { + let children = children_of(data_type); + let child_data = match children { + Some(fields) if fields.len() == data.child_data().len() => data + .child_data() + .iter() + .zip(fields.iter()) + .map(|(child, field)| relabel_data(child, field.data_type())) + .collect::>>()?, + _ => data.child_data().to_vec(), + }; + data.clone() .into_builder() .data_type(data_type.clone()) + .child_data(child_data) .build() .map_err(|e| { Error::invalid_input(format!( "a {} column cannot be read as {data_type}: {e}", - column.data_type() + data.data_type() )) - })?; - Ok(arrow_array::make_array(relabelled)) + }) } /// [`without_field_ids`] for a type rather than a schema, for the nested types a /// reconciliation builds. -fn without_field_ids_in(data_type: &DataType) -> DataType { +pub(crate) fn without_field_ids_in(data_type: &DataType) -> DataType { let one = ArrowSchema::new(vec![ArrowField::new("", data_type.clone(), true)]); without_field_ids(&one).field(0).data_type().clone() } @@ -475,3 +492,66 @@ fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) - Source::Live => Ok(Arc::new(BooleanArray::from(vec![false; rows]))), } } + +#[cfg(test)] +mod relabel_tests { + use super::*; + use arrow_array::{Int64Array, StructArray}; + use arrow_schema::Fields; + + fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { + ArrowField::new(name, data_type, true).with_metadata( + [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] + .into_iter() + .collect(), + ) + } + + /// A field id lives inside a nested column's type at every level, so the + /// relabel has to reach all of them: Arrow validates a struct's children + /// against the child types its own type declares. + #[test] + fn relabel_reaches_a_nested_child() { + let inner_stamped = stamped("b", DataType::Int64, 3); + let middle_stamped = stamped( + "inner", + DataType::Struct(Fields::from(vec![inner_stamped.clone()])), + 2, + ); + let outer_stamped = DataType::Struct(Fields::from(vec![middle_stamped.clone()])); + + let leaf = Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef; + let middle = StructArray::new( + Fields::from(vec![inner_stamped]), + vec![Arc::clone(&leaf)], + None, + ); + let outer = Arc::new(StructArray::new( + Fields::from(vec![middle_stamped]), + vec![Arc::new(middle) as ArrayRef], + None, + )) as ArrayRef; + assert_eq!(outer.data_type(), &outer_stamped); + + let plain = without_field_ids_in(&outer_stamped); + let relabelled = relabel_to(&outer, &plain).expect("relabel a nested column"); + assert_eq!(relabelled.data_type(), &plain, "every level is relabelled"); + + // The values have to survive, not just the type. + let as_struct = relabelled + .as_any() + .downcast_ref::() + .expect("a struct"); + let middle = as_struct + .column(0) + .as_any() + .downcast_ref::() + .expect("a nested struct"); + let values = middle + .column(0) + .as_any() + .downcast_ref::() + .expect("the leaf"); + assert_eq!(values.value(0), 7); + } +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index 5270e8f31e1..fc21d6fca84 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -26,7 +26,7 @@ use lance_core::is_system_column; use lance_core::{Error, Result}; use super::exec::ReconcileExec; -use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; +use crate::dataset::mem_wal::reconcile::{Plan, field_id_of, without_field_ids_in}; use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// One sealed generation, read under the table's schema. @@ -119,13 +119,10 @@ impl GenerationRead { /// generation has and the table does not. Either way the predicate belongs /// above the reconciliation, where the columns it names exist. pub(super) fn to_stored(&self, expr: &Expr) -> Option { - let pushable = expr.column_refs().iter().all(|c| { - self.stored_name(&c.name).is_some_and(|stored| { - self.stored - .field_with_name(stored) - .is_ok_and(|f| !is_reconstructed(f.data_type())) - }) - }); + let pushable = expr + .column_refs() + .iter() + .all(|c| self.answers_as_written(&c.name)); if !pushable { return None; } @@ -143,6 +140,28 @@ impl GenerationRead { .ok() } + /// Whether the generation holds `table_name` in the shape the table + /// declares, so a predicate naming it means the same thing pushed down. + /// + /// A nested column is where the two can differ without the name moving: a + /// reference names the parent (`info.a` refers to `info`) and a parent's + /// name does not move when a child is renamed. Comparing the shapes rather + /// than assuming the worst is what keeps an ordinary predicate on an + /// ordinary struct pushed down — the common case, where nothing moved. + fn answers_as_written(&self, table_name: &str) -> bool { + let Some(stored) = self.stored_name(table_name) else { + return false; + }; + let (Ok(stored), Ok(declared)) = ( + self.stored.field_with_name(stored), + self.identity.field_with_name(table_name), + ) else { + return false; + }; + // Field ids live inside a nested type, and are not part of the shape. + without_field_ids_in(stored.data_type()) == without_field_ids_in(declared.data_type()) + } + /// Bring the scan's output back to the table's names and shapes: renames /// followed, columns the generation never stored filled with nulls, nested /// columns rebuilt to the shape the table declares. @@ -186,11 +205,31 @@ impl GenerationRead { /// The schema [`Self::reconcile`] produces: the wanted columns as the table /// declares them, then whatever else the scan carries. + /// + /// This is the intermediate schema, not the public one: nullability comes + /// from the source, which is where the rows actually are. + /// + /// A generation stores every non-key column as nullable however the table + /// declares it — that is what lets a strict table hold a tombstone, whose + /// payload is null in everything but the key. A point lookup carries + /// tombstones through on purpose, so those rows have to survive + /// reconciliation. A column the generation never stored is likewise null + /// for its rows. The table's own nullability is restored at the public + /// boundary, by the canonical projection each arm passes through, once the + /// tombstones have been dropped. fn target(&self, source: &Schema) -> SchemaRef { let mut fields: Vec = self .wanted .iter() - .filter_map(|name| self.identity.field_with_name(name).ok().cloned()) + .filter_map(|name| { + let declared = self.identity.field_with_name(name).ok()?; + // Absent from the source means synthesized, so nullable. + let nullable = self + .stored_name(name) + .and_then(|stored| source.field_with_name(stored).ok()) + .is_none_or(|f| f.is_nullable()); + Some(declared.clone().with_nullable(nullable)) + }) .collect(); // A generation's own columns are not the table's, so they pass through // as the generation has them. @@ -221,18 +260,6 @@ pub(super) fn filter_above( )) } -/// A column the reconciliation rebuilds rather than takes as it stands, so its -/// stored shape is not the shape a pushed-down predicate would expect. -fn is_reconstructed(data_type: &DataType) -> bool { - match data_type { - DataType::Struct(_) => true, - DataType::List(e) | DataType::LargeList(e) | DataType::FixedSizeList(e, _) => { - is_reconstructed(e.data_type()) - } - _ => false, - } -} - /// The generation's name for each of the table's columns, by field id: a rename /// changes the name and keeps the id. pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap { @@ -463,24 +490,6 @@ mod tests { assert_eq!(read.to_stored(&col("info").is_not_null()), None); } - /// A list of structs is rebuilt too, so its stored shape is not the shape a - /// pushed-down predicate would expect. - #[test] - fn a_list_is_reconstructed_only_when_its_element_is() { - let struct_element = - DataType::Struct(Fields::from(vec![Field::new("c", DataType::Int64, true)])); - assert!(is_reconstructed(&DataType::List(Arc::new(Field::new( - "item", - struct_element, - true - ))))); - assert!(!is_reconstructed(&DataType::List(Arc::new(Field::new( - "item", - DataType::Int64, - true - ))))); - } - /// With no ids to match on, the table's own names are the only link — the /// behaviour a caller that supplies no identity schema gets. #[test] From b75b727fbbcb00607bd4f5f5dc7b3505311ed679 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 22:54:20 -0400 Subject: [PATCH 28/48] test(mem_wal): cover the nested shapes the relabel has to preserve A relabel touches every level of a nested array, so every level's structure is part of the contract: a null parent and a null child, a list's offsets with an empty and a null row, LargeList's wider offsets, FixedSizeList's width, a non-zero slice offset, and an empty column where the schema is the whole contract. Each asserts the values and the validity, not that construction succeeded. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 173 ++++++++++++++++++++ 1 file changed, 173 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index df88fb18243..5f1d0ce74b0 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -555,3 +555,176 @@ mod relabel_tests { assert_eq!(values.value(0), 7); } } + +#[cfg(test)] +mod nested_relabel_tests { + use super::*; + use arrow_array::{ + Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, StructArray, + }; + use arrow_buffer::{NullBuffer, OffsetBuffer}; + use arrow_schema::Fields; + + fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { + ArrowField::new(name, data_type, true).with_metadata( + [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] + .into_iter() + .collect(), + ) + } + + /// Relabel `column` to its own type with the field ids stripped, and check + /// that nothing but the labels moved. + fn strip_and_check(column: ArrayRef) -> ArrayRef { + let plain = without_field_ids_in(column.data_type()); + let out = relabel_to(&column, &plain).expect("relabel"); + assert_eq!(out.data_type(), &plain, "every level is relabelled"); + assert_eq!(out.len(), column.len(), "row count is preserved"); + assert_eq!( + out.null_count(), + column.null_count(), + "validity is preserved" + ); + out + } + + /// A struct whose parent is null at one row, and whose child is null at + /// another: both levels of validity have to survive the relabel. + #[test] + fn a_null_parent_and_a_null_child_both_survive() { + let child = stamped("b", DataType::Int64, 3); + let values = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; + let column = Arc::new(StructArray::new( + Fields::from(vec![child]), + vec![values], + Some(NullBuffer::from(vec![true, true, false])), + )) as ArrayRef; + + let out = strip_and_check(column); + let out = out.as_any().downcast_ref::().expect("struct"); + assert!(out.is_null(2), "the null parent stays null"); + let inner = out + .column(0) + .as_any() + .downcast_ref::() + .expect("the child"); + assert_eq!(inner.value(0), 1); + assert!(inner.is_null(1), "the null child stays null"); + } + + /// A list carries offsets and its own validity, and the element carries + /// values: an empty list, a null list and a null element in one column. + #[test] + fn a_lists_offsets_and_validity_survive() { + let element = stamped("item", DataType::Struct(Fields::from(vec![stamped("b", DataType::Int64, 4)])), 3); + let leaf = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; + let inner = Arc::new(StructArray::new( + Fields::from(vec![stamped("b", DataType::Int64, 4)]), + vec![leaf], + None, + )) as ArrayRef; + // Rows: [two elements], [], null. + let column = Arc::new(ListArray::new( + Arc::new(element), + OffsetBuffer::new(vec![0, 2, 2, 3].into()), + inner, + Some(NullBuffer::from(vec![true, true, false])), + )) as ArrayRef; + + let out = strip_and_check(column); + let out = out.as_any().downcast_ref::().expect("list"); + assert_eq!(out.value_length(0), 2, "the first row keeps two elements"); + assert_eq!(out.value_length(1), 0, "the empty list stays empty"); + assert!(out.is_null(2), "the null list stays null"); + } + + /// `LargeList` is a different offset width, and neither array downcasts to + /// the other. + #[test] + fn a_large_lists_offsets_survive() { + let element = stamped("item", DataType::Int64, 3); + let values = Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef; + let column = Arc::new(LargeListArray::new( + Arc::new(element), + OffsetBuffer::new(vec![0i64, 2, 3].into()), + values, + None, + )) as ArrayRef; + + let out = strip_and_check(column); + let out = out + .as_any() + .downcast_ref::() + .expect("a large list, not a list"); + assert_eq!(out.value_length(0), 2); + assert_eq!(out.value_length(1), 1); + } + + /// A fixed-size list's width lives in its type, so a relabel must carry it. + #[test] + fn a_fixed_size_lists_width_survives() { + let element = stamped("item", DataType::Int64, 3); + let values = Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef; + let column = + Arc::new(FixedSizeListArray::new(Arc::new(element), 2, values, None)) as ArrayRef; + + let out = strip_and_check(column); + assert!( + matches!(out.data_type(), DataType::FixedSizeList(_, 2)), + "the width is part of the type, got {:?}", + out.data_type() + ); + } + + /// A sliced array carries a non-zero offset into its buffers. Relabelling + /// must not reinterpret that as a full array. + #[test] + fn a_slice_keeps_its_offset() { + let child = stamped("b", DataType::Int64, 3); + let values = Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef; + let whole = StructArray::new(Fields::from(vec![child]), vec![values], None); + let column = Arc::new(whole.slice(2, 2)) as ArrayRef; + + let out = strip_and_check(column); + let out = out.as_any().downcast_ref::().expect("struct"); + let inner = out + .column(0) + .as_any() + .downcast_ref::() + .expect("the child"); + assert_eq!( + (0..out.len()).map(|i| inner.value(i)).collect::>(), + vec![3, 4], + "the slice reads its own rows, not the array's first ones" + ); + } + + /// An empty batch has no values to check, so the schema is the whole + /// contract. + #[test] + fn an_empty_column_is_still_relabelled() { + let child = stamped("b", DataType::Int64, 3); + let column = Arc::new(StructArray::new( + Fields::from(vec![child]), + vec![Arc::new(Int64Array::from(Vec::::new())) as ArrayRef], + None, + )) as ArrayRef; + let out = strip_and_check(column); + assert_eq!(out.len(), 0); + } + + /// Nothing to change is the fast path, and it has to return the same + /// arrays rather than a rebuilt approximation of them. + #[test] + fn a_column_already_in_the_target_shape_is_returned_as_it_stands() { + let plain = ArrowField::new("b", DataType::Int64, true); + let column = Arc::new(StructArray::new( + Fields::from(vec![plain]), + vec![Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef], + None, + )) as ArrayRef; + let out = relabel_to(&column, column.data_type()).expect("relabel"); + assert_eq!(out.data_type(), column.data_type()); + assert_eq!(out.len(), 2); + } +} From 3d9fa2bf2b642febf6ff4254903c945548f7a59e Mon Sep 17 00:00:00 2001 From: XYZhan Date: Tue, 15 Sep 2026 23:39:15 -0400 Subject: [PATCH 29/48] docs(mem_wal): state the invariant the memtable arm rests on MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It reads under the table's own names rather than resolving them, which is sound because a memtable is rebuilt with the claim — its schema is the schema the read was planned against. Routing it through GenerationRead instead was tried and dropped: the plan is non-identity purely because the storage schema carries field ids, so it would add a per-batch node to every WAL read for no change in result. Naming the invariant, and the tests that enforce it, is the part worth keeping. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 8 ++++++-- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 6 ++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 5f1d0ce74b0..51afc86e9d0 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -21,11 +21,11 @@ use std::collections::HashMap; use std::sync::Arc; +use arrow::array::ArrayData; use arrow_array::{ Array, ArrayRef, BooleanArray, FixedSizeListArray, GenericListArray, RecordBatch, RecordBatchOptions, StructArray, }; -use arrow::array::ArrayData; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::{Error, Result}; @@ -616,7 +616,11 @@ mod nested_relabel_tests { /// values: an empty list, a null list and a null element in one column. #[test] fn a_lists_offsets_and_validity_survive() { - let element = stamped("item", DataType::Struct(Fields::from(vec![stamped("b", DataType::Int64, 4)])), 3); + let element = stamped( + "item", + DataType::Struct(Fields::from(vec![stamped("b", DataType::Int64, 4)])), + 3, + ); let leaf = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; let inner = Arc::new(StructArray::new( Fields::from(vec![stamped("b", DataType::Int64, 4)]), diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index fdaf8f7348b..04ce0e97a24 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -431,6 +431,12 @@ impl LsmScanPlanner { let mut scanner = MemTableScanner::new(batch_store.clone(), index_store.clone(), schema.clone()); + // Asked for under the table's own names, which a memtable + // stores them under: a memtable is rebuilt with the claim, so + // its schema is the schema this read was planned against. That + // is the invariant this arm rests on instead of resolving one, + // and what `a_frozen_memtable_reads_the_same_before_and_after_ + // its_flush` and the lifecycle equivalence tests enforce. let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; From fc35ace61394e2690abc7711f3c3e1be2178e4ba Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 00:09:09 -0400 Subject: [PATCH 30/48] fix(mem_wal): do not truncate a search before a deferred predicate runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A predicate the generation cannot answer — one naming a column it was sealed before, or a nested path whose identity moved — runs above the reconciliation, which is after the search has already chosen its top-k. Rows that pass the predicate were being dropped behind rows that do not, so a query could come back empty while an eligible row sat just outside the cut. A compound predicate reaches this without any nested column: `added IS NULL AND flag >= n` over a generation added to since it sealed. Those arms now do not cut. The vector arm ranks every row the generation holds and the full-text arm takes no limit, leaving the filter and then the union's own top-k to decide. Only a generation whose predicate could not be translated pays for it; the general scan already withheld its limit for the same reason. Pushdown eligibility also now compares a column's stored and declared types with their field ids rather than stripping them first. Two nested children can share a name and a type and still be different columns — one dropped and one added back — and a predicate pushed down against the retired child would filter on values the reconciliation above replaces with nulls. --- rust/lance/src/dataset/mem_wal/scanner/fts_search.rs | 12 ++++++++---- rust/lance/src/dataset/mem_wal/scanner/generation.rs | 10 +++++++--- .../src/dataset/mem_wal/scanner/vector_search.rs | 11 +++++++++++ 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index e0a5298120e..0eb5697fbad 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -981,10 +981,14 @@ impl LsmFtsSearchPlanner { scanner.prefilter(true); } let mut bound_query = query.clone().with_column(stored_column)?; - if let Some(limit) = limit { - bound_query = bound_query.limit(Some(limit as i64)); - } else { - bound_query = bound_query.limit(None); + // A predicate that could not be pushed down runs above the + // reconciliation, which is after the query has taken its top-k + // by score. Cutting first would drop rows that pass the + // predicate behind rows that do not, so a generation with a + // deferred predicate does not cut. + match limit.filter(|_| above.is_none()) { + Some(limit) => bound_query = bound_query.limit(Some(limit as i64)), + None => bound_query = bound_query.limit(None), } scanner.full_text_search(bound_query)?; let reconciled = generation.reconcile(scanner.create_plan().await?)?; diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index fc21d6fca84..e9f582b0563 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -26,7 +26,7 @@ use lance_core::is_system_column; use lance_core::{Error, Result}; use super::exec::ReconcileExec; -use crate::dataset::mem_wal::reconcile::{Plan, field_id_of, without_field_ids_in}; +use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// One sealed generation, read under the table's schema. @@ -158,8 +158,12 @@ impl GenerationRead { ) else { return false; }; - // Field ids live inside a nested type, and are not part of the shape. - without_field_ids_in(stored.data_type()) == without_field_ids_in(declared.data_type()) + // Compared with their field ids, not just their shapes. A nested child + // that was dropped and added back under the same name and type is a + // different column wearing the old one's shape, and pushing a predicate + // down against it would filter on the retired child's values while the + // reconciliation above synthesizes nulls for the new one. + stored.data_type() == declared.data_type() } /// Bring the scan's output back to the table's names and shapes: renames diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 316829cb455..9f13ea777e9 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -554,6 +554,17 @@ impl LsmVectorSearchPlanner { } // No `with_row_id/address`: per-source IDs would collide with base. let query_arr = single_query_array(query_vector); + // A predicate that could not be pushed down runs above the + // reconciliation, which is after the search has chosen its + // top-k. Cutting to `k` first would drop rows that pass the + // predicate behind rows that do not, so this arm does not cut: + // it ranks everything it holds and lets the filter, and then + // the union's own top-k, decide. Only a generation the + // predicate cannot be translated against pays for this. + let k = match above { + None => k, + Some(_) => dataset.count_rows(None).await?.max(1), + }; scanner.nearest(&vector_column, query_arr.as_ref(), k)?; scanner.distance_range(self.distance_range.0, self.distance_range.1); scanner.nprobes(nprobes); From ccd229abc3af96b436ce641d2b0641343e738d2e Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 00:53:18 -0400 Subject: [PATCH 31/48] fix(mem_wal): refuse tightening a column to non-null on a WAL-backed table The check runs against the base table, and a write admitted into the WAL while it runs is not there to be checked. Draining first narrows the window without closing it: a drain waits for the generations it sealed, and writes keep being admitted into the next one, so a null can be acknowledged after the check and before the commit. The row then sits in a table whose schema forbids it, and every compaction that would fold it into base fails from then on. Closing it properly needs a table-scoped write barrier held across the validation and the commit, which the WAL has no way to ask for. Refusing is what the retype case already does for the same reason, so the two now share one refusal. A table without a MemWAL tightens exactly as before. --- rust/lance/src/dataset/schema_evolution.rs | 50 ++++++++++++++++++---- 1 file changed, 42 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 669dc848bad..4d1efbf8fbf 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -484,15 +484,42 @@ pub(super) async fn add_columns( /// /// Takes the decision as a bool rather than the alterations themselves: a /// reference to them held across an await would have to be `Sync`. -async fn reject_cast_on_mem_wal(dataset: &Dataset, casts: bool) -> Result<()> { - if !casts || dataset.mem_wal_index_details().await?.is_none() { +/// What `alter_columns` refuses on a table with a MemWAL, and why. +/// +/// Both are alterations whose meaning depends on rows the commit cannot see. +enum Unsupported { + /// A cast takes a new field id, which rows still in the WAL cannot be + /// matched to. + Retype, + /// Tightening a column to non-null is validated against base, and a write + /// admitted into the WAL while that runs is not in base to be validated. + /// Draining first does not close it: the drain waits for the generations it + /// sealed, and writes keep being admitted into the next one, so a null can + /// be acknowledged after the check and before the commit. The row is then + /// in a table whose schema forbids it, and the compaction that would fold + /// it into base fails from then on. + Tightening, +} + +async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) -> Result<()> { + let Some(unsupported) = unsupported else { + return Ok(()); + }; + if dataset.mem_wal_index_details().await?.is_none() { return Ok(()); } - Err(Error::invalid_input( - "cannot change a column's type on a table with a MemWAL attached: a cast takes a \ - new field id, which rows still in the WAL cannot be matched to. Drop the MemWAL, \ - or add a column of the new type and backfill it.", - )) + Err(Error::invalid_input(match unsupported { + Unsupported::Retype => { + "cannot change a column's type on a table with a MemWAL attached: a cast takes a \ + new field id, which rows still in the WAL cannot be matched to. Drop the MemWAL, \ + or add a column of the new type and backfill it." + } + Unsupported::Tightening => { + "cannot make a column non-nullable on a table with a MemWAL attached: the check \ + runs against the base table, and a write admitted into the WAL while it runs is \ + not there to be checked. Drop the MemWAL first." + } + })) } async fn cleanup_new_column_data_files(fragments: &[FileFragment], new_fragments: &[Fragment]) { @@ -753,7 +780,14 @@ pub(super) async fn alter_columns( dataset: &mut Dataset, alterations: &[ColumnAlteration], ) -> Result<()> { - reject_cast_on_mem_wal(dataset, alterations.iter().any(|a| a.data_type.is_some())).await?; + let unsupported = if alterations.iter().any(|a| a.data_type.is_some()) { + Some(Unsupported::Retype) + } else if alterations.iter().any(|a| a.nullable == Some(false)) { + Some(Unsupported::Tightening) + } else { + None + }; + reject_on_mem_wal(dataset, unsupported).await?; // Validate referenced columns exist and enforce NOT NULL when tightening // a column from nullable to non-nullable. From dfa75eb1d0a3ee7057d15b7b96644a99ed1ace0a Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 01:52:41 -0400 Subject: [PATCH 32/48] fix(mem_wal): refuse renaming a nested field onto a sibling's name Compaction matches a struct's children by name. A pair of renames that exchange two children's names leaves both names on both sides, so the merge cannot tell which child is which and writes each one's values under the other's -- silently, and the generation is then trimmed. Resolving nested children by field id in the compactor is more machinery than this case is worth: a nested child cannot be added to an existing struct at all (the add is rejected), a nested drop and an ordinary nested rename both come through correctly, and only a rename onto a name a sibling currently holds is ambiguous. That one is refused, which is what the retype and tightening cases already do for their own undecidable inputs. Top-level renames are untouched: those are resolved by field id, so a pair of them can exchange names safely. --- rust/lance/src/dataset/schema_evolution.rs | 39 ++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 4d1efbf8fbf..23d6a1c4a12 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -491,6 +491,13 @@ enum Unsupported { /// A cast takes a new field id, which rows still in the WAL cannot be /// matched to. Retype, + /// Renaming a nested child onto a name one of its siblings currently holds. + /// Compaction matches a struct's children by name, so both sides would + /// carry the same two names and the merge could not tell which child is + /// which -- it would write each one's values under the other's name. + /// Resolving nested children by id there is more machinery than this case + /// is worth; a rename that does not collide with a sibling is unaffected. + SiblingNameReuse, /// Tightening a column to non-null is validated against base, and a write /// admitted into the WAL while that runs is not in base to be validated. /// Draining first does not close it: the drain waits for the generations it @@ -501,6 +508,27 @@ enum Unsupported { Tightening, } +/// Whether `alteration` renames a nested field onto a name one of its siblings +/// currently holds. See [`Unsupported::SiblingNameReuse`]. +/// +/// Top-level renames are exempt: compaction resolves those by field id, so a +/// pair of them can exchange names safely. +fn renames_onto_a_sibling(dataset: &Dataset, alteration: &ColumnAlteration) -> bool { + let Some(rename) = &alteration.rename else { + return false; + }; + let Some((parent, child)) = alteration.path.rsplit_once('.') else { + return false; + }; + let Some(parent) = dataset.schema().field(parent) else { + return false; + }; + parent + .children + .iter() + .any(|sibling| sibling.name != child && sibling.name == *rename) +} + async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) -> Result<()> { let Some(unsupported) = unsupported else { return Ok(()); @@ -519,6 +547,12 @@ async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) runs against the base table, and a write admitted into the WAL while it runs is \ not there to be checked. Drop the MemWAL first." } + Unsupported::SiblingNameReuse => { + "cannot rename a nested field onto a name one of its siblings holds on a table \ + with a MemWAL attached: rows still in the WAL would have the two children \ + matched by name and their values exchanged. Rename the sibling out of the way \ + first, or drop the MemWAL." + } })) } @@ -784,6 +818,11 @@ pub(super) async fn alter_columns( Some(Unsupported::Retype) } else if alterations.iter().any(|a| a.nullable == Some(false)) { Some(Unsupported::Tightening) + } else if alterations + .iter() + .any(|a| renames_onto_a_sibling(dataset, a)) + { + Some(Unsupported::SiblingNameReuse) } else { None }; From 60fbdb370a4b8af16353ff3d338b218f29e4031a Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 02:19:31 -0400 Subject: [PATCH 33/48] refactor(mem_wal): one test module for the relabel, and narrow what it exports The two modules were added a round apart and both test `relabel_to`, each with its own copy of the same field-stamping helper. `stored_names` was public to the scanner while only this module calls it. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 53 +++++++------------ .../src/dataset/mem_wal/scanner/generation.rs | 19 +++---- 2 files changed, 26 insertions(+), 46 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 51afc86e9d0..bd68160f729 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -496,7 +496,10 @@ fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) - #[cfg(test)] mod relabel_tests { use super::*; - use arrow_array::{Int64Array, StructArray}; + use arrow_array::{ + Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, StructArray, + }; + use arrow_buffer::{NullBuffer, OffsetBuffer}; use arrow_schema::Fields; fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { @@ -507,6 +510,21 @@ mod relabel_tests { ) } + /// Relabel `column` to its own type with the field ids stripped, and check + /// that nothing but the labels moved. + fn strip_and_check(column: ArrayRef) -> ArrayRef { + let plain = without_field_ids_in(column.data_type()); + let out = relabel_to(&column, &plain).expect("relabel"); + assert_eq!(out.data_type(), &plain, "every level is relabelled"); + assert_eq!(out.len(), column.len(), "row count is preserved"); + assert_eq!( + out.null_count(), + column.null_count(), + "validity is preserved" + ); + out + } + /// A field id lives inside a nested column's type at every level, so the /// relabel has to reach all of them: Arrow validates a struct's children /// against the child types its own type declares. @@ -554,39 +572,6 @@ mod relabel_tests { .expect("the leaf"); assert_eq!(values.value(0), 7); } -} - -#[cfg(test)] -mod nested_relabel_tests { - use super::*; - use arrow_array::{ - Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, StructArray, - }; - use arrow_buffer::{NullBuffer, OffsetBuffer}; - use arrow_schema::Fields; - - fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { - ArrowField::new(name, data_type, true).with_metadata( - [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] - .into_iter() - .collect(), - ) - } - - /// Relabel `column` to its own type with the field ids stripped, and check - /// that nothing but the labels moved. - fn strip_and_check(column: ArrayRef) -> ArrayRef { - let plain = without_field_ids_in(column.data_type()); - let out = relabel_to(&column, &plain).expect("relabel"); - assert_eq!(out.data_type(), &plain, "every level is relabelled"); - assert_eq!(out.len(), column.len(), "row count is preserved"); - assert_eq!( - out.null_count(), - column.null_count(), - "validity is preserved" - ); - out - } /// A struct whose parent is null at one row, and whose child is null at /// another: both levels of validity have to survive the relabel. diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index e9f582b0563..442fbdf31c1 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -210,17 +210,12 @@ impl GenerationRead { /// The schema [`Self::reconcile`] produces: the wanted columns as the table /// declares them, then whatever else the scan carries. /// - /// This is the intermediate schema, not the public one: nullability comes - /// from the source, which is where the rows actually are. - /// - /// A generation stores every non-key column as nullable however the table - /// declares it — that is what lets a strict table hold a tombstone, whose - /// payload is null in everything but the key. A point lookup carries - /// tombstones through on purpose, so those rows have to survive - /// reconciliation. A column the generation never stored is likewise null - /// for its rows. The table's own nullability is restored at the public - /// boundary, by the canonical projection each arm passes through, once the - /// tombstones have been dropped. + /// Nullability comes from the source, not the table: a generation stores + /// every non-key column as nullable, which is what lets a strict table hold + /// a tombstone. A point lookup carries tombstones through on purpose, so + /// those rows have to survive this. The table's own nullability is restored + /// by the canonical projection each arm passes through, after the + /// tombstones are dropped. fn target(&self, source: &Schema) -> SchemaRef { let mut fields: Vec = self .wanted @@ -266,7 +261,7 @@ pub(super) fn filter_above( /// The generation's name for each of the table's columns, by field id: a rename /// changes the name and keeps the id. -pub(super) fn stored_names(stored: &Schema, table: &Schema) -> HashMap { +fn stored_names(stored: &Schema, table: &Schema) -> HashMap { let by_id: HashMap = table .fields() .iter() From 9bc5453e8240354bdccbfa31d3e1199f7d022b71 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 02:39:08 -0400 Subject: [PATCH 34/48] docs(mem_wal): describe the memtable and tightening contracts in Lance's terms Both comments leaned on a deployment's vocabulary -- a claim being rebuilt, a drain waiting on generations -- and one pointed at a test that lives in a different repository. Restated against what this crate has: the schema a memtable's writer holds, and the window a flush leaves open. --- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 12 ++++++------ rust/lance/src/dataset/schema_evolution.rs | 13 ++++++------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 04ce0e97a24..db2b95921bf 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -431,12 +431,12 @@ impl LsmScanPlanner { let mut scanner = MemTableScanner::new(batch_store.clone(), index_store.clone(), schema.clone()); - // Asked for under the table's own names, which a memtable - // stores them under: a memtable is rebuilt with the claim, so - // its schema is the schema this read was planned against. That - // is the invariant this arm rests on instead of resolving one, - // and what `a_frozen_memtable_reads_the_same_before_and_after_ - // its_flush` and the lifecycle equivalence tests enforce. + // Asked for under the table's own names, which is what a + // memtable stores them under: a memtable is created from the + // schema its writer holds, so a reader planning against that + // same schema needs no resolution. A caller pairing a memtable + // with a newer schema is outside this contract -- pass the + // memtable its own schema, or reopen the writer. let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 23d6a1c4a12..8bb0f8d10a9 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -498,13 +498,12 @@ enum Unsupported { /// Resolving nested children by id there is more machinery than this case /// is worth; a rename that does not collide with a sibling is unaffected. SiblingNameReuse, - /// Tightening a column to non-null is validated against base, and a write - /// admitted into the WAL while that runs is not in base to be validated. - /// Draining first does not close it: the drain waits for the generations it - /// sealed, and writes keep being admitted into the next one, so a null can - /// be acknowledged after the check and before the commit. The row is then - /// in a table whose schema forbids it, and the compaction that would fold - /// it into base fails from then on. + /// Tightening a column to non-null is validated against the committed + /// fragments, and a row still in the MemWAL is not among them. Flushing + /// first does not close the window: a flush covers the generations open + /// when it starts, and writes keep arriving into the next one, so a null + /// can be accepted after the check and before the commit. That row is then + /// in a table whose schema forbids it, and every later merge of it fails. Tightening, } From c3e59b2f63ae59a20c42e558702785cdf9f2395f Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 02:44:54 -0400 Subject: [PATCH 35/48] style(mem_wal): satisfy the workspace's denied lints `redundant_pub_crate` and `use_self` are deny-level here and CI runs clippy with `-D warnings`, so these would have failed it. --- rust/lance/src/dataset/mem_wal/reconcile.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index bd68160f729..a53a3d68dc9 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -33,7 +33,7 @@ use lance_core::{Error, Result}; use super::TOMBSTONE; /// The lance field id an Arrow field carries, if it carries one. -pub(crate) fn field_id_of(field: &ArrowField) -> Option { +pub(super) fn field_id_of(field: &ArrowField) -> Option { field .metadata() .get(LANCE_FIELD_ID_KEY) @@ -52,7 +52,7 @@ pub(crate) fn field_id_of(field: &ArrowField) -> Option { /// id-carrying storage schema and the table's plain one describe the same /// values as two different types. Only the labels differ, so this relabels the /// array rather than converting it. -pub(crate) fn relabel_to(column: &ArrayRef, data_type: &DataType) -> Result { +pub(super) fn relabel_to(column: &ArrayRef, data_type: &DataType) -> Result { if column.data_type() == data_type { return Ok(column.clone()); } @@ -88,12 +88,12 @@ fn relabel_data(data: &ArrayData, data_type: &DataType) -> Result { /// [`without_field_ids`] for a type rather than a schema, for the nested types a /// reconciliation builds. -pub(crate) fn without_field_ids_in(data_type: &DataType) -> DataType { +pub(super) fn without_field_ids_in(data_type: &DataType) -> DataType { let one = ArrowSchema::new(vec![ArrowField::new("", data_type.clone(), true)]); without_field_ids(&one).field(0).data_type().clone() } -pub(crate) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { +pub(super) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { fn strip(field: &ArrowField) -> ArrowField { let mut metadata = field.metadata().clone(); metadata.remove(LANCE_FIELD_ID_KEY); @@ -129,7 +129,7 @@ enum Source { Take(usize), /// The source column at this index, whose struct children need their own /// resolution. - Nested(usize, Vec, DataType), + Nested(usize, Vec, DataType), /// The source does not have this column: rows written before it existed /// hold no value for it. Null(DataType), From c02751cfa7634e8b6db31008857344f90d4dcf58 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:01:15 -0400 Subject: [PATCH 36/48] fix(mem_wal): keep the generation arms' futures inside the recursion limit Each arm's plan build is boxed at the call site, as the point-lookup arms already are: a generation resolves its own schema before scanning, and leaving those futures inlined pushes the `Send` proof past rustc's limit for the callers stacked above them. Also drops an unused test import that `-D warnings` rejects. --- rust/lance/src/dataset/mem_wal/scanner/fts_search.rs | 2 +- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 6 +++++- rust/lance/src/dataset/mem_wal/scanner/vector_search.rs | 4 ++-- rust/lance/src/dataset/mem_wal/write.rs | 2 +- 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 0eb5697fbad..16d7feefc17 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -991,7 +991,7 @@ impl LsmFtsSearchPlanner { None => bound_query = bound_query.limit(None), } scanner.full_text_search(bound_query)?; - let reconciled = generation.reconcile(scanner.create_plan().await?)?; + let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; match above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index db2b95921bf..c348047a36b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -414,7 +414,11 @@ impl LsmScanPlanner { scanner.limit(Some(fetch as i64), None)?; } - let reconciled = generation.reconcile(scanner.create_plan().await?)?; + // Boxed at the call site, as the point-lookup arms are: the + // generation's own planning nests deeply enough that leaving + // this future inlined pushes the `Send` proof past rustc's + // recursion limit for callers stacked above it. + let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; match above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 9f13ea777e9..a582b29e2b5 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -563,7 +563,7 @@ impl LsmVectorSearchPlanner { // predicate cannot be translated against pays for this. let k = match above { None => k, - Some(_) => dataset.count_rows(None).await?.max(1), + Some(_) => Box::pin(dataset.count_rows(None)).await?.max(1), }; scanner.nearest(&vector_column, query_arr.as_ref(), k)?; scanner.distance_range(self.distance_range.0, self.distance_range.1); @@ -573,7 +573,7 @@ impl LsmVectorSearchPlanner { scanner.ef(ef); } scanner.fast_search(); - let reconciled = generation.reconcile(scanner.create_plan().await?)?; + let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; match above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 1456fc9cd1a..fb7aa654765 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -4580,7 +4580,7 @@ pub fn new_shared_stats() -> SharedWriteStats { mod tests { use super::*; use crate::dataset::mem_wal::test_util::failing_memory_store; - use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, Int64Array, StringArray}; + use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, StringArray}; use arrow_schema::Field as ArrowField; use arrow_schema::{DataType, Field}; use lance_core::FenceReason; From 8c29607eb9fafa2d6ca39376806ffa0597fe7218 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:07:21 -0400 Subject: [PATCH 37/48] fix(mem_wal): keep ReconcileExec out of the public surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing outside the crate constructs it, and its constructor takes a `Plan`, which is crate-internal — so documenting it publicly meant a public item linking to a private one. --- rust/lance/src/dataset/mem_wal/scanner/exec.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/scanner/exec.rs index a353b10c27e..46de19e87d3 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec.rs @@ -31,5 +31,5 @@ pub use pk::{ validate_pk_types, }; pub use pk_block_filter::PkBlockFilterExec; -pub use reconcile::ReconcileExec; +pub(crate) use reconcile::ReconcileExec; pub use schema_relabel::SchemaRelabelExec; From c96d8831b26c8bc3663ffc633ed04cb83aab58f4 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:11:28 -0400 Subject: [PATCH 38/48] docs(mem_wal): say why the search arms box their plan build The scan planner's arm carries the reason; the two search arms had the same `Box::pin` with nothing explaining it. --- rust/lance/src/dataset/mem_wal/scanner/fts_search.rs | 3 +++ rust/lance/src/dataset/mem_wal/scanner/vector_search.rs | 3 +++ 2 files changed, 6 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 16d7feefc17..439ebb1744f 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -991,6 +991,9 @@ impl LsmFtsSearchPlanner { None => bound_query = bound_query.limit(None), } scanner.full_text_search(bound_query)?; + // Boxed for the reason the scan planner's arm gives: a + // generation resolves its own schema before scanning, and + // the inlined future is too deep for the `Send` proof. let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; match above { Some(expr) => filter_above(reconciled, expr), diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index a582b29e2b5..3e633cb7375 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -573,6 +573,9 @@ impl LsmVectorSearchPlanner { scanner.ef(ef); } scanner.fast_search(); + // Boxed for the reason the scan planner's arm gives: a + // generation resolves its own schema before scanning, and + // the inlined future is too deep for the `Send` proof. let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; match above { Some(expr) => filter_above(reconciled, expr), From 9534bf4abb605b5433713a4749dc993fb35a6414 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:15:23 -0400 Subject: [PATCH 39/48] fix(mem_wal): box the scan planner's per-source arm The vector and full-text planners already box theirs. An arm resolves a generation's schema before scanning, and holding that across the scan's own future leaves the whole chain's `Send` proof past rustc's recursion limit for the callers stacked above it. --- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index c348047a36b..ec477bd0747 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -192,9 +192,11 @@ impl LsmScanPlanner { (Some(n), false, false) => Some(n), _ => None, }; - let scan = self - .build_source_scan(&source, projection, filter, fetch) - .await?; + // Boxed per arm, as the vector and full-text planners box theirs: + // an arm resolves a generation's schema before scanning, and + // leaving its future inlined here puts the whole scan chain's + // `Send` proof past rustc's recursion limit. + let scan = Box::pin(self.build_source_scan(&source, projection, filter, fetch)).await?; // Drop cross-generation stale rows (PKs superseded by a newer gen). // Plain scans refill exactly, so keep the approximate-search From 6949a192bea37ea3bf182f7eeadf4a61ad3377ad Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:19:02 -0400 Subject: [PATCH 40/48] fix(mem_wal): type-erase each source arm's future Boxing alone leaves the concrete future in place, so the `Send` proof still recurses through it; a trait object is where that recursion stops. Both the scan and point-lookup arms resolve a generation's schema before scanning, which nests deeply enough that the callers stacked above them exceeded rustc's recursion limit. --- rust/lance/src/dataset/mem_wal/scanner/planner.rs | 12 +++++++----- .../src/dataset/mem_wal/scanner/point_lookup.rs | 10 +++++++--- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index ec477bd0747..afd9c58d083 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -192,11 +192,13 @@ impl LsmScanPlanner { (Some(n), false, false) => Some(n), _ => None, }; - // Boxed per arm, as the vector and full-text planners box theirs: - // an arm resolves a generation's schema before scanning, and - // leaving its future inlined here puts the whole scan chain's - // `Send` proof past rustc's recursion limit. - let scan = Box::pin(self.build_source_scan(&source, projection, filter, fetch)).await?; + // Type-erased, not merely boxed: the `Send` proof recurses + // through a boxed future's concrete type but stops at a trait + // object. An arm resolves a generation's schema before it + // scans, which nests deeply enough to need that. + let arm: futures::future::BoxFuture<'_, Result>> = + Box::pin(self.build_source_scan(&source, projection, filter, fetch)); + let scan = arm.await?; // Drop cross-generation stale rows (PKs superseded by a newer gen). // Plain scans refill exactly, so keep the approximate-search diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index 0fcc1a59ebb..681356f5c37 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -275,9 +275,13 @@ impl LsmPointLookupPlanner { for source in sources { let generation = source.generation().as_u64(); - let scan = self - .build_source_scan(&source, projection, &filter_expr) - .await?; + // Type-erased, not merely boxed: the `Send` proof recurses + // through a boxed future's concrete type but stops at a trait + // object. An arm resolves a generation's schema before it + // scans, which nests deeply enough to need that. + let arm: futures::future::BoxFuture<'_, Result>> = + Box::pin(self.build_source_scan(&source, projection, &filter_expr)); + let scan = arm.await?; // Data is stored in reverse order, so first match is newest let limited: Arc = Arc::new(GlobalLimitExec::new(scan, 0, Some(1))); From 9922e487273310eac22d96a50d764fb506d3894c Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 04:33:01 -0400 Subject: [PATCH 41/48] fix(mem_wal): type-erase the search arms, and drop an out-of-scope doc link The `Send` proof reaches the vector and full-text arms through the same per-source fan-out the scan and point-lookup arms use, so those are erased too. `LANCE_FIELD_ID_KEY` is not in scope where it was linked. --- .../src/dataset/mem_wal/scanner/fts_search.rs | 21 +++++++++++-------- .../dataset/mem_wal/scanner/vector_search.rs | 21 ++++++++++++------- rust/lance/src/dataset/mem_wal/write.rs | 2 +- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 439ebb1744f..abf58f4ec9b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -804,17 +804,20 @@ impl LsmFtsSearchPlanner { (source, is_active, blocked, fetch_limit) }) .collect(); + // Type-erased for the reason the vector planner's arm gives. let built = futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, fetch_limit)| { - Box::pin(self.build_source_plan( - source, - column, - &query, - *fetch_limit, - projection, - index_params.as_ref(), - &target_schema, - )) + let arm: futures::future::BoxFuture<'_, Result>> = + Box::pin(self.build_source_plan( + source, + column, + &query, + *fetch_limit, + projection, + index_params.as_ref(), + &target_schema, + )); + arm })) .await?; diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 3e633cb7375..86999f2fed4 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -328,16 +328,21 @@ impl LsmVectorSearchPlanner { (source, is_base, is_active, blocked, fetch_k) }) .collect(); + // Type-erased, not merely boxed: the `Send` proof recurses through a + // boxed future's concrete type but stops at a trait object, and an arm + // resolves a generation's schema before it searches. let built = futures::future::try_join_all(arm_inputs.iter().map( |(source, is_base, _, _, fetch_k)| { - Box::pin(self.build_knn_plan( - source, - query_vector, - *fetch_k, - nprobes, - projection, - *is_base && refine_base, - )) + let arm: futures::future::BoxFuture<'_, Result>> = + Box::pin(self.build_knn_plan( + source, + query_vector, + *fetch_k, + nprobes, + projection, + *is_base && refine_base, + )); + arm }, )) .await?; diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index fb7aa654765..be9592b35fd 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -2098,7 +2098,7 @@ impl ShardWriter { /// The `base_path` should come from `ObjectStore::from_uri()` to ensure /// WAL files are written inside the dataset directory. /// - /// `schema` carrying each field's id under [`LANCE_FIELD_ID_KEY`] in its + /// `schema` carrying each field's id under `lance:field_id` in its /// field metadata is what lets a replayed entry be matched to a column that /// has since been renamed. Without them, a replayed entry is matched by /// name, and a renamed column reads as null for every row the memtable From 0237dd43d4d19a8e4f7a461158933c82940384af Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 05:05:39 -0400 Subject: [PATCH 42/48] fix(mem_wal): carry the recursive contract through maps and quoted paths A map is a nested container like any other -- its entries struct holds the key and value as children with ids of their own -- so it now takes part in the id stamping, the traversal, the id stripping, the relabel and the rebuild, with its offsets, validity and sortedness carried through. The sibling-rename guard split its path on the last dot, which lands inside a quoted component: a field whose own name contains dots named the wrong parent, and the guard admitted the rename it exists to refuse. It resolves the path instead. --- rust/lance/src/dataset/mem_wal.rs | 5 + rust/lance/src/dataset/mem_wal/reconcile.rs | 119 ++++++++++++++++++ .../src/dataset/mem_wal/scanner/generation.rs | 5 + .../src/dataset/mem_wal/scanner/planner.rs | 6 +- rust/lance/src/dataset/schema_evolution.rs | 50 +++++++- 5 files changed, 179 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index d9c88802e35..2f6c563a05f 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -186,6 +186,11 @@ fn stamp_field_id(field: &ArrowField, among: &[Field]) -> ArrowField { let element = stamp_field_id(element, &source.children); field.with_data_type(DataType::FixedSizeList(Arc::new(element), size)) } + DataType::Map(entries, sorted) => { + let sorted = *sorted; + let entries = stamp_field_id(entries, &source.children); + field.with_data_type(DataType::Map(Arc::new(entries), sorted)) + } _ => field, } } diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index a53a3d68dc9..7dec4938096 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -115,6 +115,12 @@ pub(super) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { .clone() .with_data_type(DataType::FixedSizeList(Arc::new(strip(element)), size)) } + DataType::Map(entries, sorted) => { + let sorted = *sorted; + field + .clone() + .with_data_type(DataType::Map(Arc::new(strip(entries)), sorted)) + } _ => field, } } @@ -335,6 +341,7 @@ fn is_nested(data_type: &DataType) -> bool { | DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + | DataType::Map(_, _) ) } @@ -348,6 +355,9 @@ fn children_of(data_type: &DataType) -> Option { Some(vec![element.as_ref().clone()].into()) } DataType::FixedSizeList(element, _) => Some(vec![element.as_ref().clone()].into()), + // A map's child is its entries struct, which carries the key and value + // as children of its own. + DataType::Map(entries, _) => Some(vec![entries.as_ref().clone()].into()), _ => None, } } @@ -432,6 +442,41 @@ fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) - Source::Nested(i, children, to @ DataType::LargeList(_)) => { rebuild_list::(&columns[*i], &children[0], to, name) } + // A map is its entries struct behind offsets. The sortedness flag is + // part of the type, so it comes from the target with the rest of it. + Source::Nested(i, children, to @ DataType::Map(_, sorted)) => { + let map = columns[*i] + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::invalid_input(format!("column `{name}` is not a map")))?; + let Some(entries) = children_of(to).and_then(|c| c.first().cloned()) else { + unreachable!("a map target has an entries field"); + }; + let stored_entries: ArrayRef = Arc::new(map.entries().clone()); + let rebuilt = take_column( + &children[0], + std::slice::from_ref(&stored_entries), + map.entries().len(), + entries.name(), + )?; + let rebuilt = rebuilt + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input(format!("map column `{name}` entries are not a struct")) + })? + .clone(); + Ok(Arc::new( + arrow_array::MapArray::try_new( + entries, + map.offsets().clone(), + rebuilt, + map.nulls().cloned(), + *sorted, + ) + .map_err(|e| Error::invalid_input(format!("rebuild map column `{name}`: {e}")))?, + )) + } // A fixed-size list is rebuilt the same way, keeping its width. Source::Nested(i, children, to @ DataType::FixedSizeList(_, _)) => { let DataType::FixedSizeList(_, size) = to else { @@ -573,6 +618,80 @@ mod relabel_tests { assert_eq!(values.value(0), 7); } + /// A map is a nested container like any other: its entries carry field ids + /// of their own, so a rename inside one has to resolve by id rather than + /// read as a changed type. + #[test] + fn a_renamed_field_inside_a_map_resolves_by_id() { + let entry = |value: &str| { + ArrowField::new( + "entries", + DataType::Struct(Fields::from(vec![ + stamped("key", DataType::Int64, 2), + stamped(value, DataType::Int64, 3), + ])), + false, + ) + }; + let keys = Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef; + let values = Arc::new(Int64Array::from(vec![10, 20])) as ArrayRef; + let entries = StructArray::new( + Fields::from(vec![ + stamped("key", DataType::Int64, 2), + stamped("old", DataType::Int64, 3), + ]), + vec![keys, values], + None, + ); + let column = Arc::new( + arrow_array::MapArray::try_new( + Arc::new(entry("old")), + arrow_buffer::OffsetBuffer::new(vec![0, 2].into()), + entries, + None, + false, + ) + .unwrap(), + ) as ArrayRef; + + let source = ArrowSchema::new(vec![stamped( + "m", + DataType::Map(Arc::new(entry("old")), false), + 1, + )]); + let target: SchemaRef = Arc::new(ArrowSchema::new(vec![stamped( + "m", + DataType::Map(Arc::new(entry("new")), false), + 1, + )])); + let batch = RecordBatch::try_new(Arc::new(source.clone()), vec![column]).unwrap(); + + let plan = Plan::resolve(&source, &target, &[]).expect("a rename inside a map resolves"); + let out = plan.apply(&batch).expect("apply"); + let map = out + .column(0) + .as_any() + .downcast_ref::() + .expect("a map"); + assert_eq!( + map.entries().column_names(), + vec!["key", "new"], + "the renamed entry arrives under its new name" + ); + let vals = map + .entries() + .column(1) + .as_any() + .downcast_ref::() + .expect("values"); + assert_eq!( + (vals.value(0), vals.value(1)), + (10, 20), + "carrying its values" + ); + assert_eq!(map.value_length(0), 2, "and its offsets"); + } + /// A struct whose parent is null at one row, and whose child is null at /// another: both levels of validity have to survive the relabel. #[test] diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index 442fbdf31c1..4b65bb19bc9 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -335,6 +335,11 @@ fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { size, )) } + (DataType::Map(entries, sorted), DataType::Map(source_entries, _)) => { + let sorted = *sorted; + let one: Fields = vec![source_entries.as_ref().clone()].into(); + field.with_data_type(DataType::Map(Arc::new(restore(entries, &one)), sorted)) + } _ => field, } } diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index afd9c58d083..3bf21b4aabc 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -442,9 +442,9 @@ impl LsmScanPlanner { // Asked for under the table's own names, which is what a // memtable stores them under: a memtable is created from the // schema its writer holds, so a reader planning against that - // same schema needs no resolution. A caller pairing a memtable - // with a newer schema is outside this contract -- pass the - // memtable its own schema, or reopen the writer. + // same schema needs no resolution. Pairing a memtable with a + // schema it was not created from is outside this contract -- + // pass the memtable its own schema, or reopen the writer. let cols = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 8bb0f8d10a9..58ae3fee226 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -516,16 +516,19 @@ fn renames_onto_a_sibling(dataset: &Dataset, alteration: &ColumnAlteration) -> b let Some(rename) = &alteration.rename else { return false; }; - let Some((parent, child)) = alteration.path.rsplit_once('.') else { + // Resolved rather than split on `.`: a field name may contain dots, in + // which case the path quotes it, and splitting would name the wrong parent + // and admit exactly the rename this refuses. + let Some(chain) = dataset.schema().resolve(&alteration.path) else { return false; }; - let Some(parent) = dataset.schema().field(parent) else { + let [.., parent, child] = chain.as_slice() else { return false; }; parent .children .iter() - .any(|sibling| sibling.name != child && sibling.name == *rename) + .any(|sibling| sibling.name != child.name && sibling.name == *rename) } async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) -> Result<()> { @@ -4631,4 +4634,45 @@ mod test { Ok(()) } + + /// A field name may contain dots, in which case the path quotes it. + /// Splitting on the last dot names the wrong parent, and the sibling check + /// then admits the rename it exists to refuse. + #[tokio::test] + async fn a_quoted_path_resolves_to_its_real_parent() -> Result<()> { + let child = ArrowField::new("child.with.dot", DataType::Int32, true); + let sibling = ArrowField::new("sibling", DataType::Int32, true); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "parent", + DataType::Struct(ArrowFields::from(vec![child.clone(), sibling.clone()])), + true, + )])); + let parent = StructArray::new( + ArrowFields::from(vec![child, sibling]), + vec![ + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + Arc::new(Int32Array::from(vec![2])) as ArrayRef, + ], + None, + ); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(parent)])?; + let test_dir = TempStrDir::default(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = Dataset::write(reader, &test_dir, None).await?; + + let onto_sibling = + ColumnAlteration::new("parent.`child.with.dot`".into()).rename("sibling".into()); + assert!( + renames_onto_a_sibling(&dataset, &onto_sibling), + "renaming onto a sibling's name must be seen" + ); + + let onto_free = + ColumnAlteration::new("parent.`child.with.dot`".into()).rename("free".into()); + assert!( + !renames_onto_a_sibling(&dataset, &onto_free), + "a name no sibling holds is not a collision" + ); + Ok(()) + } } From 8f89f5a6cdc5d1a555230e80f1a4ea514aead063 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 10:23:37 -0400 Subject: [PATCH 43/48] docs(mem_wal): describe the skipped index in Lance's own terms --- rust/lance/src/dataset/mem_wal/api.rs | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index 9387acfbf32..3155c3dca46 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -751,14 +751,9 @@ async fn build_index_configs( .next(); // An index the maintained set names and the dataset does not have: - // dropped outright, or carried away with the column it covered. The set - // is fixed at `initialize_mem_wal` and cannot be edited afterwards, so - // refusing here refuses the claim -- and a table whose claim cannot be - // built serves no reads at all, for a condition that costs only the - // fresh tier's copy of one index. - // - // Serve without it. The base index is gone for everyone; the - // fresh tier simply has nothing to keep in step with. + // dropped outright, or carried away with the column it covered. Serve + // without it -- the base index is gone for everyone, so the fresh tier + // has nothing to keep in step with. See `MissingIndex::Skip`. let Some(index_meta) = index_meta else { if on_missing == OnMissingIndex::Reject { return Err(Error::invalid_input(format!( From c0b1498261ff9c010a0640a2c66ef09c33f46bcc Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 11:02:19 -0400 Subject: [PATCH 44/48] refactor(mem_wal): name the two schemas for whose they are --- rust/lance/src/dataset/mem_wal.rs | 12 +- rust/lance/src/dataset/mem_wal/api.rs | 25 +- rust/lance/src/dataset/mem_wal/reconcile.rs | 145 ++++++----- .../src/dataset/mem_wal/scanner/builder.rs | 12 +- .../src/dataset/mem_wal/scanner/fts_search.rs | 27 +-- .../src/dataset/mem_wal/scanner/generation.rs | 227 ++++++++++++------ .../src/dataset/mem_wal/scanner/planner.rs | 19 +- .../dataset/mem_wal/scanner/point_lookup.rs | 4 +- .../dataset/mem_wal/scanner/vector_search.rs | 26 +- rust/lance/src/dataset/mem_wal/write.rs | 58 +++-- rust/lance/src/dataset/schema_evolution.rs | 146 +++++++---- 11 files changed, 414 insertions(+), 287 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 2f6c563a05f..af7df745dc5 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -116,12 +116,6 @@ pub fn relax_non_pk_nullability( )) } -/// Extend the logical schema with the trailing `_tombstone` column — the -/// intermediate [`relax_non_pk_nullability`] widens into the storage schema. -/// -/// Idempotent: a schema that already carries `_tombstone` (a reopen/replay -/// path) is returned unchanged. Schema-level metadata and per-field metadata -/// (e.g. the `lance-schema:unenforced-primary-key` marker) are preserved. /// The schema's Arrow form, with each field's id carried in its metadata. /// /// `From<&Field> for ArrowField` drops the id, which leaves everything @@ -195,6 +189,12 @@ fn stamp_field_id(field: &ArrowField, among: &[Field]) -> ArrowField { } } +/// Extend the logical schema with the trailing `_tombstone` column — the +/// intermediate [`relax_non_pk_nullability`] widens into the storage schema. +/// +/// Idempotent: a schema that already carries `_tombstone` (a reopen/replay +/// path) is returned unchanged. Schema-level metadata and per-field metadata +/// (e.g. the `lance-schema:unenforced-primary-key` marker) are preserved. pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { if base.column_with_name(TOMBSTONE).is_some() { return Arc::new(base.clone()); diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index 3155c3dca46..e0635fb24a2 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -8,11 +8,16 @@ //! //! # Limitations //! -//! MemWAL does not track dataset changes made after it is initialized: dropping -//! or replacing a maintained index, or projecting away its column, leaves -//! `maintained_indexes` naming something the writer cannot build. A change that -//! races the initialization commit lands the same way. Both surface as a failing -//! `mem_wal_writer`; handling them is follow-up work. +//! MemWAL does not track dataset changes made after it is initialized. The set +//! named by `maintained_indexes` is fixed, so an index created later is not +//! maintained over the fresh tier; it covers a row once that row reaches the +//! base table. +//! +//! An index the set names but the dataset no longer has -- dropped, replaced, or +//! carried away with the column it covered -- is skipped when a shard opens, so +//! the table keeps serving without the fresh tier's copy of it. Naming one that +//! does not exist is still rejected at initialization, which is the last moment +//! it can be corrected. use std::collections::HashMap; use std::sync::Arc; @@ -714,11 +719,6 @@ impl DatasetMemWalExt for Dataset { } } -/// Build the in-memory index configurations for `index_names`. -/// -/// Shared by [`DatasetMemWalExt::mem_wal_writer`] and -/// [`validate_maintained_indexes`], so a set that validates is one the writer -/// can build. /// Whether an index the set names but the dataset does not have is fatal. #[derive(Clone, Copy, PartialEq, Eq)] enum OnMissingIndex { @@ -732,6 +732,11 @@ enum OnMissingIndex { Skip, } +/// Build the in-memory index configurations for `index_names`. +/// +/// Shared by [`DatasetMemWalExt::mem_wal_writer`] and +/// [`validate_maintained_indexes`], so a set that validates is one the writer +/// can build. async fn build_index_configs( dataset: &Dataset, index_names: &[String], diff --git a/rust/lance/src/dataset/mem_wal/reconcile.rs b/rust/lance/src/dataset/mem_wal/reconcile.rs index 7dec4938096..c3914acc9de 100644 --- a/rust/lance/src/dataset/mem_wal/reconcile.rs +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -33,6 +33,10 @@ use lance_core::{Error, Result}; use super::TOMBSTONE; /// The lance field id an Arrow field carries, if it carries one. +/// +/// Lance writes `-1` for a field it has not assigned an id to yet. Treating +/// that as an id would pair every unassigned column with every other, so a +/// negative id counts as none at all. pub(super) fn field_id_of(field: &ArrowField) -> Option { field .metadata() @@ -41,17 +45,12 @@ pub(super) fn field_id_of(field: &ArrowField) -> Option { .filter(|id| *id >= 0) } -/// `schema` without the field ids, for comparing against what a caller sends. -/// -/// Ids belong to the stored schema, where identity has to survive a rename. A -/// caller's batch carries none, and Arrow compares a struct's children by their -/// full field — metadata included — so a stamped schema would reject it. -/// One stored column under the type the caller declared. +/// `column` under `data_type`, which differs from its own only in the field +/// ids it carries. /// -/// A field id lives inside a nested column's own Arrow type, so the memtable's -/// id-carrying storage schema and the table's plain one describe the same -/// values as two different types. Only the labels differ, so this relabels the -/// array rather than converting it. +/// The two describe the same values as two different Arrow types, because the +/// ids sit inside the type. Only the labels differ, so this relabels the array +/// rather than converting it. pub(super) fn relabel_to(column: &ArrayRef, data_type: &DataType) -> Result { if column.data_type() == data_type { return Ok(column.clone()); @@ -93,11 +92,21 @@ pub(super) fn without_field_ids_in(data_type: &DataType) -> DataType { without_field_ids(&one).field(0).data_type().clone() } +/// `field` without its Lance field id, keeping the rest of its metadata. +pub(super) fn without_field_id(field: &ArrowField) -> ArrowField { + let mut field = field.clone(); + field.metadata_mut().remove(LANCE_FIELD_ID_KEY); + field +} + +/// `schema` without the field ids, for comparing against what a caller sends. +/// +/// Ids belong to the stored schema, where identity has to survive a rename. A +/// caller's batch carries none, and Arrow compares a struct's children by their +/// full field — metadata included — so a stamped schema would reject it. pub(super) fn without_field_ids(schema: &ArrowSchema) -> ArrowSchema { fn strip(field: &ArrowField) -> ArrowField { - let mut metadata = field.metadata().clone(); - metadata.remove(LANCE_FIELD_ID_KEY); - let field = field.clone().with_metadata(metadata); + let field = without_field_id(field); match field.data_type() { DataType::Struct(children) => { let children: Vec = children.iter().map(|c| strip(c)).collect(); @@ -155,12 +164,11 @@ pub struct Plan { impl Plan { /// The same plan emitting the table's plain Arrow schema. /// - /// Resolution needs field ids on the target — a rename keeps the id and - /// moves the name — but a reader is handed the table's schema, which does - /// not carry them. They live inside a nested column's own type, so a struct + /// Resolution needs ids on the target, but a reader is handed the table's + /// plain schema. Since the ids sit inside a nested column's type, a struct /// built to the id-carrying target is a different Arrow type from the one - /// the caller declared. Only replay, which writes back into the memtable's - /// id-carrying storage schema, keeps them. + /// the caller declared. Only replay keeps them, writing back into the + /// memtable's id-carrying storage schema. pub(crate) fn emitting_plain_schema(mut self) -> Self { fn strip(source: &mut Source) { match source { @@ -189,7 +197,7 @@ impl Plan { // An id match takes its source column; a name match may then only take // one nothing has claimed. A rename frees a name for another column to // use, and it is the id that says which column is really which. - let claimed = claimed_by_id(source, target); + let claimed = claimed_by_id(source.fields(), target.fields()); let sources = target .fields() .iter() @@ -232,23 +240,6 @@ impl Plan { } } -/// Source columns an id match has taken, which a name match may not take again. -fn claimed_by_id(source: &ArrowSchema, target: &SchemaRef) -> Vec { - let by_id: HashMap = source - .fields() - .iter() - .enumerate() - .filter_map(|(i, f)| field_id_of(f).map(|id| (id, i))) - .collect(); - let mut claimed = vec![false; source.fields().len()]; - for field in target.fields() { - if let Some(i) = field_id_of(field).and_then(|id| by_id.get(&id)) { - claimed[*i] = true; - } - } - claimed -} - fn resolve_field( field: &ArrowField, source_fields: &arrow_schema::Fields, @@ -261,14 +252,9 @@ fn resolve_field( .iter() .position(|f| field_id_of(f) == Some(id)) }); - // Identity is the field id where both sides carry one. A name is not: a - // rename moves the name and leaves the id, so a source column of the same - // name under a *different* id is a different column — one dropped and - // another added under its name, whose values the table no longer has. - // - // A name match is right only where identity is absent: a batch a caller has - // just handed in carries no ids, and neither does a schema supplied by a - // caller who has none to give. + // Ids first. The name fallback is asymmetric on purpose: it fires for a + // target field carrying no id, and for one whose id no source field has -- + // an id-bearing target against an unstamped source still matches by name. let by_name = || { source_fields .iter() @@ -335,14 +321,7 @@ fn resolve_field( /// Whether this type carries its children inside its own type, so an array of /// it has to be rebuilt rather than taken as it stands. fn is_nested(data_type: &DataType) -> bool { - matches!( - data_type, - DataType::Struct(_) - | DataType::List(_) - | DataType::LargeList(_) - | DataType::FixedSizeList(_, _) - | DataType::Map(_, _) - ) + children_of(data_type).is_some() } /// The child fields of a nested type, if it has them. @@ -376,14 +355,15 @@ fn resolve_children(source: &ArrowField, field: &ArrowField) -> Result Vec { +/// Source columns an id match has taken, which a name match may not take again. +fn claimed_by_id(source: &arrow_schema::Fields, target: &arrow_schema::Fields) -> Vec { let by_id: HashMap = source .iter() .enumerate() @@ -429,6 +409,7 @@ fn rebuild_list( )) } +/// One stored column under the type the caller declared. fn take_column(source: &Source, columns: &[ArrayRef], rows: usize, name: &str) -> Result { match source { Source::Take(i) => Ok(Arc::clone(&columns[*i])), @@ -547,6 +528,57 @@ mod relabel_tests { use arrow_buffer::{NullBuffer, OffsetBuffer}; use arrow_schema::Fields; + /// A nested column resolves by the field ids Lance puts on its children. + /// If Lance starts giving children to a type [`is_nested`] does not know, + /// those ids go unrestored and the column falls back to matching by name -- + /// the failure this module exists to prevent. So every type Lance gives + /// children to has to be one we recurse into. + #[test] + fn every_type_lance_gives_children_to_is_one_we_recurse_into() { + let item = || Arc::new(ArrowField::new("item", DataType::Int64, true)); + let a_struct = || Fields::from(vec![ArrowField::new("a", DataType::Int64, true)]); + let entries = Arc::new(ArrowField::new( + "entries", + DataType::Struct(Fields::from(vec![ + ArrowField::new("key", DataType::Int64, false), + ArrowField::new("value", DataType::Int64, true), + ])), + false, + )); + let candidates = [ + DataType::Int64, + DataType::Struct(a_struct()), + DataType::List(item()), + DataType::LargeList(item()), + DataType::FixedSizeList(item(), 2), + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Struct(a_struct()), true)), + 2, + ), + DataType::Map(entries, false), + DataType::ListView(item()), + DataType::LargeListView(item()), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::RunEndEncoded( + Arc::new(ArrowField::new("run_ends", DataType::Int32, false)), + item(), + ), + ]; + for data_type in candidates { + let field = ArrowField::new("c", data_type.clone(), true); + // A type Lance refuses outright can never reach a MemWAL. + let Ok(lance) = lance_core::datatypes::Field::try_from(&field) else { + continue; + }; + if !lance.children.is_empty() { + assert!( + is_nested(&data_type), + "Lance gives {data_type:?} children, so reconcile must recurse into it" + ); + } + } + } + fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { ArrowField::new(name, data_type, true).with_metadata( [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] @@ -570,9 +602,8 @@ mod relabel_tests { out } - /// A field id lives inside a nested column's type at every level, so the - /// relabel has to reach all of them: Arrow validates a struct's children - /// against the child types its own type declares. + /// Arrow validates a struct against the child types its own type declares, + /// so a relabel that stops at the outer level produces a rejected array. #[test] fn relabel_reaches_a_nested_child() { let inner_stamped = stamped("b", DataType::Int64, 3); diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs index 6794c97b325..b78615f9512 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs @@ -348,12 +348,8 @@ impl LsmScanner { self } - /// Register a shard's in-memory memtables (active + frozen-awaiting- - /// flush) captured atomically by `ShardWriter::in_memory_memtable_refs`. - /// The read path's entry point — closes the concurrent-read-vs-flush - /// hole by carrying frozen-undrained generations into the scan. - /// Supply `schema` with each field's id, so a sealed generation's columns - /// are resolved to the table's by id rather than by name. A rename keeps + /// Supply [`Self::schema`] with each field's id, so a sealed generation's + /// columns resolve to the table's by id rather than by name. A rename keeps /// the id and moves the name, so without this a renamed column reads as /// absent. Built with /// [`arrow_schema_with_field_ids`](crate::dataset::mem_wal::arrow_schema_with_field_ids). @@ -364,6 +360,10 @@ impl LsmScanner { self } + /// Register a shard's in-memory memtables (active + frozen-awaiting- + /// flush) captured atomically by `ShardWriter::in_memory_memtable_refs`. + /// The read path's entry point — closes the concurrent-read-vs-flush + /// hole by carrying frozen-undrained generations into the scan. pub fn with_in_memory_memtables( mut self, shard_id: Uuid, diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index abf58f4ec9b..1c8dcb874d1 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -470,7 +470,6 @@ impl LsmFtsSearchPlanner { self } - /// Set the session used to open SSTables. /// The table's schema carrying each field's id, which is what resolves a /// generation's stored columns to the table's. /// @@ -481,6 +480,7 @@ impl LsmFtsSearchPlanner { self } + /// Set the session used to open SSTables. pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); self @@ -699,8 +699,8 @@ impl LsmFtsSearchPlanner { // offers no granularity. let generation = GenerationRead::new( dataset.schema(), - Arc::clone(&self.identity_schema), - self.pk_columns.clone(), + &self.identity_schema, + &self.pk_columns, Vec::new(), ); match generation.stored_name(column) { @@ -946,12 +946,12 @@ impl LsmFtsSearchPlanner { let mut scanner = dataset.scan(); // Asked of this generation under its own names: a rename moved // the table's name while the file still holds the old one. - let wanted = self.fts_scanner_projection(projection); + let asked_for = self.fts_scanner_projection(projection); let mut generation = GenerationRead::new( dataset.schema(), - Arc::clone(&self.identity_schema), - self.pk_columns.clone(), - wanted, + &self.identity_schema, + &self.pk_columns, + asked_for, ); // The index is on this generation's own column, under the name // it had when the generation was sealed. @@ -966,16 +966,7 @@ impl LsmFtsSearchPlanner { // The BM25 top-k has already run by then, so the arm can come // back short — which is right: those rows have no value for a // column sealed before it existed. - let stored_filter = self - .filter - .as_ref() - .and_then(|expr| generation.to_stored(expr)); - let above = self.filter.as_ref().filter(|_| stored_filter.is_none()); - if let Some(expr) = above { - for column in expr.column_refs() { - generation.also_produce(&column.name); - } - } + let (stored_filter, above) = generation.split_filter(self.filter.as_ref()); scanner.project(&generation.stored_projection())?; if let Some(ref stored) = stored_filter { // See the base arm: `prefilter(true)` makes this a true @@ -998,7 +989,7 @@ impl LsmFtsSearchPlanner { // generation resolves its own schema before scanning, and // the inlined future is too deep for the `Send` proof. let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; - match above { + match &above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), } diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index 4b65bb19bc9..d0cf1aa045e 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -21,12 +21,11 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::filter::FilterExec; use datafusion::prelude::Expr; use datafusion_physical_expr::create_physical_expr; -use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::is_system_column; use lance_core::{Error, Result}; use super::exec::ReconcileExec; -use crate::dataset::mem_wal::reconcile::{Plan, field_id_of}; +use crate::dataset::mem_wal::reconcile::{Plan, field_id_of, without_field_id}; use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// One sealed generation, read under the table's schema. @@ -39,58 +38,65 @@ use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// ([`Self::reconcile`]). pub(super) struct GenerationRead { /// The generation's own schema, carrying its field ids. - stored: Schema, - /// The generation's name for a column → the table's name for it. Columns - /// the table no longer has are absent. + stored_schema: Schema, + /// The generation's name for a column → the table's name for it, matched + /// by field id. Carried as names because everything downstream is + /// name-addressed: a scan projection takes names, and so do the column + /// references in a DataFusion predicate. Columns the table no longer has + /// are absent. names: HashMap, - /// The table's schema, carrying field ids. - identity: SchemaRef, + /// The table's schema, also carrying field ids. + table_schema: SchemaRef, pk_columns: Vec, - /// Table names this read produces, in order. - wanted: Vec, + /// The columns this read produces, in order, each under the name the + /// table gives it. Not only what the caller asked for -- + /// [`Self::also_produce`] adds what a deferred predicate needs. + projection: Vec, } impl GenerationRead { - /// `wanted` is what the caller asks for, in the table's names. + /// `projection` is what the caller asks for, under the table's names. pub(super) fn new( dataset_schema: &lance_core::datatypes::Schema, - identity: SchemaRef, - pk_columns: Vec, - wanted: Vec, + table_schema: &SchemaRef, + pk_columns: &[String], + projection: Vec, ) -> Self { - let stored = arrow_schema_with_field_ids(dataset_schema); - let names = stored_names(&stored, &identity); + let stored_schema = arrow_schema_with_field_ids(dataset_schema); + let names = stored_names(&stored_schema, table_schema); + let (table_schema, pk_columns) = (Arc::clone(table_schema), pk_columns.to_vec()); Self { - stored, + stored_schema, names, - identity, + table_schema, pk_columns, - wanted, + projection, } } - /// The generation's name for one of the table's columns. - pub(super) fn stored_name(&self, table_name: &str) -> Option<&str> { + /// The generation's name for `column`. `column` is the table's name for it. + pub(super) fn stored_name(&self, column: &str) -> Option<&str> { self.names .iter() - .find(|(_, table)| *table == table_name) - .map(|(stored, _)| stored.as_str()) + .find(|(_, in_table)| *in_table == column) + .map(|(in_generation, _)| in_generation.as_str()) } /// Also produce `column`, which the caller needs even though it did not ask /// for it — a predicate that runs after reconciliation reads its columns /// from this scan. pub(super) fn also_produce(&mut self, column: &str) { - if !self.wanted.iter().any(|w| w == column) { - self.wanted.push(column.to_string()); + if !self.projection.iter().any(|w| w == column) { + self.projection.push(column.to_string()); } } - /// What to project from the file: the wanted columns under the names it has. + /// What to project from the file: the projected columns under the names it + /// has. /// A column it never stored is dropped here and filled in by /// [`Self::reconcile`]. pub(super) fn stored_projection(&self) -> Vec<&str> { - self.wanted + self.projection .iter() .filter_map(|name| { self.stored_name(name) @@ -104,7 +110,7 @@ impl GenerationRead { /// asked for, or not at all. fn stored_system_column(&self, name: &str) -> Option<&str> { (is_system_column(name) || name == TOMBSTONE) - .then(|| self.stored.field_with_name(name).ok()) + .then(|| self.stored_schema.field_with_name(name).ok()) .flatten() .map(|f| f.name().as_str()) } @@ -112,17 +118,18 @@ impl GenerationRead { /// `expr` with each column reference moved to the name this generation /// stores it under, so it can be pushed into the generation's own scan. /// - /// `None` when it cannot be: a column the generation does not store, or a - /// nested one. A nested reference names the parent (`info.a` refers to - /// `info`) and a parent's name does not move when a child is renamed, so a - /// pushed-down predicate would be evaluated against child names this - /// generation has and the table does not. Either way the predicate belongs - /// above the reconciliation, where the columns it names exist. + /// `None` when a referenced column is one this generation does not store, + /// or stores under a different shape. A nested reference names the parent + /// (`info.a` refers to `info`), and a parent's name does not move when a + /// child is renamed — so the parent's whole type is compared, not its name. + /// A nested column nothing moved inside is pushed down like any other. + /// Otherwise the predicate belongs above the reconciliation, where the + /// columns it names exist. pub(super) fn to_stored(&self, expr: &Expr) -> Option { let pushable = expr .column_refs() .iter() - .all(|c| self.answers_as_written(&c.name)); + .all(|c| self.stored_as_declared(&c.name)); if !pushable { return None; } @@ -130,8 +137,7 @@ impl GenerationRead { .transform(|e| match e { Expr::Column(mut c) => { // `stored_name` is total over the refs, checked above. - let stored = self.stored_name(&c.name).expect("checked").to_string(); - c.name = stored; + c.name = self.stored_name(&c.name).expect("checked").to_string(); Ok(Transformed::yes(Expr::Column(c))) } other => Ok(Transformed::no(other)), @@ -140,30 +146,48 @@ impl GenerationRead { .ok() } - /// Whether the generation holds `table_name` in the shape the table - /// declares, so a predicate naming it means the same thing pushed down. + /// Split `filter` into the part this generation can answer under its own + /// names and the part that has to run above the reconciliation. + /// + /// The deferred part reads its columns from this scan, so they are added to + /// what the scan produces whether the caller asked for them or not. At most + /// one of the two is `Some`: a predicate is pushed whole or deferred whole. + pub(super) fn split_filter(&mut self, filter: Option<&Expr>) -> (Option, Option) { + let pushed = filter.and_then(|expr| self.to_stored(expr)); + let deferred = filter.filter(|_| pushed.is_none()).cloned(); + if let Some(expr) = &deferred { + for column in expr.column_refs() { + self.also_produce(&column.name); + } + } + (pushed, deferred) + } + + /// Whether the generation stores `column` exactly as the table declares + /// it, so a predicate naming it means the same thing pushed down. /// /// A nested column is where the two can differ without the name moving: a /// reference names the parent (`info.a` refers to `info`) and a parent's /// name does not move when a child is renamed. Comparing the shapes rather /// than assuming the worst is what keeps an ordinary predicate on an /// ordinary struct pushed down — the common case, where nothing moved. - fn answers_as_written(&self, table_name: &str) -> bool { - let Some(stored) = self.stored_name(table_name) else { + fn stored_as_declared(&self, column: &str) -> bool { + let Some(stored_column) = self.stored_name(column) else { return false; }; - let (Ok(stored), Ok(declared)) = ( - self.stored.field_with_name(stored), - self.identity.field_with_name(table_name), + let (Ok(stored_field), Ok(declared)) = ( + self.stored_schema.field_with_name(stored_column), + self.table_schema.field_with_name(column), ) else { return false; }; - // Compared with their field ids, not just their shapes. A nested child - // that was dropped and added back under the same name and type is a - // different column wearing the old one's shape, and pushing a predicate - // down against it would filter on the retired child's values while the - // reconciliation above synthesizes nulls for the new one. - stored.data_type() == declared.data_type() + // The top-level ids matched already: `stored_name` resolved through a + // map keyed by them. This compares the children, whose ids live inside + // the parent's own type -- so a child dropped and added back under the + // same name and type is caught as the different column it is, rather + // than filtering on the retired child's values while the reconciliation + // above synthesizes nulls for the new one. + stored_field.data_type() == declared.data_type() } /// Bring the scan's output back to the table's names and shapes: renames @@ -174,7 +198,8 @@ impl GenerationRead { /// those pass through untouched, as does anything the generation has that /// the table does not. pub(super) fn reconcile(&self, scan: Arc) -> Result> { - let source = self.only_the_tables_ids(with_ids_from(&scan.schema(), &self.stored)); + let source = + self.with_only_table_field_ids(with_ids_from(&scan.schema(), &self.stored_schema)); let target = self.target(&source); let plan = Plan::resolve(&source, &target, &self.pk_columns)?.emitting_plain_schema(); if plan.is_identity() { @@ -183,31 +208,27 @@ impl GenerationRead { Ok(Arc::new(ReconcileExec::new(scan, Arc::new(plan)))) } - /// `source` with the field ids of everything that is not one of the - /// table's columns removed. + /// `source` keeping a field id only where the column is one of the + /// table's, and stripping it everywhere else. /// /// A generation numbers its own columns in its own schema — `_tombstone`, /// and anything it holds that the table has since dropped — so those ids /// collide with whatever the table gave those numbers. Left in place, a /// column added to the table resolves to whichever of them happens to share /// its id. - fn only_the_tables_ids(&self, source: Schema) -> Schema { + fn with_only_table_field_ids(&self, source: Schema) -> Schema { let fields: Vec = source .fields() .iter() .map(|field| match self.names.contains_key(field.name()) { true => field.as_ref().clone(), - false => { - let mut metadata = field.metadata().clone(); - metadata.remove(LANCE_FIELD_ID_KEY); - field.as_ref().clone().with_metadata(metadata) - } + false => without_field_id(field), }) .collect(); Schema::new_with_metadata(fields, source.metadata().clone()) } - /// The schema [`Self::reconcile`] produces: the wanted columns as the table + /// The schema [`Self::reconcile`] produces: the projected columns as the table /// declares them, then whatever else the scan carries. /// /// Nullability comes from the source, not the table: a generation stores @@ -218,14 +239,14 @@ impl GenerationRead { /// tombstones are dropped. fn target(&self, source: &Schema) -> SchemaRef { let mut fields: Vec = self - .wanted + .projection .iter() .filter_map(|name| { - let declared = self.identity.field_with_name(name).ok()?; + let declared = self.table_schema.field_with_name(name).ok()?; // Absent from the source means synthesized, so nullable. let nullable = self .stored_name(name) - .and_then(|stored| source.field_with_name(stored).ok()) + .and_then(|stored_column| source.field_with_name(stored_column).ok()) .is_none_or(|f| f.is_nullable()); Some(declared.clone().with_nullable(nullable)) }) @@ -250,7 +271,7 @@ pub(super) fn filter_above( ) -> Result> { let schema = plan.schema(); let df_schema = DFSchema::try_from(schema.as_ref().clone()) - .map_err(|e| Error::internal(format!("filter schema: {e}")))?; + .map_err(|e| Error::internal(format!("build a filter schema for `{expr}`: {e}")))?; let props = ExecutionProps::new(); let physical = create_physical_expr(expr, &df_schema, &props) .map_err(|e| Error::internal(format!("plan filter `{expr}`: {e}")))?; @@ -259,25 +280,26 @@ pub(super) fn filter_above( )) } -/// The generation's name for each of the table's columns, by field id: a rename +/// Each column the generation stores, keyed by the name it stores it under, +/// mapped to the table's name for it. Paired by field id, since a rename /// changes the name and keeps the id. -fn stored_names(stored: &Schema, table: &Schema) -> HashMap { - let by_id: HashMap = table +fn stored_names(stored_schema: &Schema, table_schema: &Schema) -> HashMap { + let by_id: HashMap = table_schema .fields() .iter() .filter_map(|f| field_id_of(f).map(|id| (id, f.name().as_str()))) .collect(); // A caller that supplies no ids leaves only names to match on. if by_id.is_empty() { - return stored + return stored_schema .fields() .iter() .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) - .filter(|f| table.field_with_name(f.name()).is_ok()) + .filter(|f| table_schema.field_with_name(f.name()).is_ok()) .map(|f| (f.name().clone(), f.name().clone())) .collect(); } - stored + stored_schema .fields() .iter() // A generation's own columns are numbered in its own schema, so their @@ -295,7 +317,7 @@ fn stored_names(stored: &Schema, table: &Schema) -> HashMap { /// Put back the field ids a scan's output schema drops, so the reconciliation /// can resolve its columns by id. -fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { +fn with_ids_from(schema: &Schema, stored_schema: &Schema) -> Schema { fn restore(field: &Field, among: &Fields) -> Field { let Some(source) = among.iter().find(|f| f.name() == field.name()) else { return field.clone(); @@ -346,7 +368,7 @@ fn with_ids_from(schema: &Schema, stored: &Schema) -> Schema { let fields: Vec = schema .fields() .iter() - .map(|field| restore(field, stored.fields())) + .map(|field| restore(field, stored_schema.fields())) .collect(); Schema::new_with_metadata(fields, schema.metadata().clone()) } @@ -356,6 +378,7 @@ mod tests { use super::*; use arrow_schema::Fields; use datafusion::prelude::{col, lit}; + use lance_core::datatypes::LANCE_FIELD_ID_KEY; use lance_core::datatypes::Schema as LanceSchema; /// An Arrow field carrying a Lance field id, as a generation's schema and @@ -374,13 +397,17 @@ mod tests { /// `GenerationRead` resolves against a generation's *Lance* schema, which /// is where the stored ids come from. - fn generation(stored: SchemaRef, table: SchemaRef, wanted: &[&str]) -> GenerationRead { - let lance = LanceSchema::try_from(stored.as_ref()).expect("a lance schema"); + fn generation( + stored_schema: SchemaRef, + table_schema: SchemaRef, + projection: &[&str], + ) -> GenerationRead { + let lance = LanceSchema::try_from(stored_schema.as_ref()).expect("a lance schema"); GenerationRead::new( &lance, - table, - vec!["id".to_string()], - wanted.iter().map(|s| s.to_string()).collect(), + &table_schema, + &["id".to_string()], + projection.iter().map(|s| s.to_string()).collect(), ) } @@ -478,7 +505,7 @@ mod tests { /// when a child is renamed — so pushing it down would evaluate it against /// child names the table does not have. #[test] - fn a_predicate_on_a_nested_column_is_never_pushed_down() { + fn a_predicate_on_a_struct_whose_child_was_renamed_is_not_pushed_down() { let nested = |child: &str| { with_id( "info", @@ -494,20 +521,60 @@ mod tests { assert_eq!(read.to_stored(&col("info").is_not_null()), None); } + /// The common case: a struct nothing moved is pushed down like any other + /// column. Deferring these was the bug -- a search takes its top-k first, + /// so a predicate applied afterwards loses a lower-ranked row that should + /// have won. + #[test] + fn a_predicate_on_a_struct_that_did_not_move_is_pushed_down() { + let unmoved = with_id( + "info", + DataType::Struct(Fields::from(vec![with_id("c", DataType::Int64, 2)])), + 1, + ); + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0), unmoved.clone()]), + schema(vec![with_id("id", DataType::Int64, 0), unmoved]), + &["id", "info"], + ); + let expr = col("info").is_not_null(); + assert_eq!(read.to_stored(&expr), Some(expr)); + } + + /// Same name, same type, different field id: a child dropped and added + /// back is a different column wearing the old one's shape, so a predicate + /// on it must not reach the retired values. + #[test] + fn a_predicate_on_a_struct_whose_child_was_replaced_is_not_pushed_down() { + let child = |id: i32| { + with_id( + "info", + DataType::Struct(Fields::from(vec![with_id("c", DataType::Int64, id)])), + 1, + ) + }; + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0), child(2)]), + schema(vec![with_id("id", DataType::Int64, 0), child(9)]), + &["id", "info"], + ); + assert_eq!(read.to_stored(&col("info").is_not_null()), None); + } + /// With no ids to match on, the table's own names are the only link — the /// behaviour a caller that supplies no identity schema gets. #[test] fn a_table_without_ids_matches_by_name() { - let stored = Schema::new(vec![ + let stored_schema = Schema::new(vec![ with_id("id", DataType::Int64, 0), with_id("value", DataType::Int64, 1), with_id(TOMBSTONE, DataType::Boolean, 2), ]); - let table = Schema::new(vec![ + let table_schema = Schema::new(vec![ Field::new("id", DataType::Int64, true), Field::new("value", DataType::Int64, true), ]); - let names = stored_names(&stored, &table); + let names = stored_names(&stored_schema, &table_schema); assert_eq!(names.get("value"), Some(&"value".to_string())); assert_eq!(names.get(TOMBSTONE), None, "not one of the table's columns"); } diff --git a/rust/lance/src/dataset/mem_wal/scanner/planner.rs b/rust/lance/src/dataset/mem_wal/scanner/planner.rs index 3bf21b4aabc..061e72971b8 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -326,6 +326,7 @@ impl LsmScanPlanner { Arc::new(Schema::new(fields)) } + /// Build scan plan for a single data source. async fn build_source_scan( &self, source: &LsmDataSource, @@ -372,13 +373,13 @@ impl LsmScanPlanner { // Asked of this generation under its own names, so an older // file is only asked for columns it has. A column it never had // is filled in after the scan. - let wanted = + let asked_for = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); let mut generation = GenerationRead::new( dataset.schema(), - Arc::clone(&self.identity_schema), - self.pk_columns.clone(), - wanted, + &self.identity_schema, + &self.pk_columns, + asked_for, ); // A predicate the generation can answer is pushed into its scan // under the names it has. One it cannot — because it names a @@ -386,13 +387,7 @@ impl LsmScanPlanner { // the reconciliation instead, reading its columns from this // scan, so they have to be in it whether the caller asked or // not. - let stored_filter = filter.and_then(|expr| generation.to_stored(expr)); - let above = filter.filter(|_| stored_filter.is_none()); - if let Some(expr) = above { - for column in expr.column_refs() { - generation.also_produce(&column.name); - } - } + let (stored_filter, above) = generation.split_filter(filter); scanner.project(&generation.stored_projection())?; scanner.with_row_address(); @@ -423,7 +418,7 @@ impl LsmScanPlanner { // this future inlined pushes the `Send` proof past rustc's // recursion limit for callers stacked above it. let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; - match above { + match &above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), } diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index 681356f5c37..2bd0c3a910d 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -706,8 +706,8 @@ impl LsmPointLookupPlanner { let cols = cols_with_tombstone(&cols, dataset.schema().field(TOMBSTONE).is_some()); let generation = GenerationRead::new( dataset.schema(), - Arc::clone(&self.identity_schema), - self.pk_columns.clone(), + &self.identity_schema, + &self.pk_columns, cols, ); // Every generation stores every primary key column — a key diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs index 86999f2fed4..61ae36bd854 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -179,7 +179,6 @@ impl LsmVectorSearchPlanner { self } - /// Set the session used to open SSTables. /// The table's schema carrying each field's id, which is what resolves a /// generation's stored columns to the table's. /// @@ -190,6 +189,7 @@ impl LsmVectorSearchPlanner { self } + /// Set the session used to open SSTables. pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); self @@ -519,13 +519,13 @@ impl LsmVectorSearchPlanner { // the table's name while the file still holds the old one, so // projecting the table's names would ask for a column that is // not there. - let wanted = + let asked_for = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); let mut generation = GenerationRead::new( dataset.schema(), - Arc::clone(&self.identity_schema), - self.pk_columns.clone(), - wanted, + &self.identity_schema, + &self.pk_columns, + asked_for, ); // The index is on this generation's own column, under the name // it had when the generation was sealed. @@ -537,19 +537,7 @@ impl LsmVectorSearchPlanner { let vector_column = vector_column.to_string(); // A predicate this generation cannot answer as written runs // above the reconciliation, where the columns it names exist. - // The search's own top-k has already run by then, so the arm - // can come back short — which is right: those rows have no - // value for a column sealed before it existed. - let stored_filter = self - .filter - .as_ref() - .and_then(|expr| generation.to_stored(expr)); - let above = self.filter.as_ref().filter(|_| stored_filter.is_none()); - if let Some(expr) = above { - for column in expr.column_refs() { - generation.also_produce(&column.name); - } - } + let (stored_filter, above) = generation.split_filter(self.filter.as_ref()); scanner.project(&generation.stored_projection())?; if let Some(ref stored) = stored_filter { // See the base arm: `prefilter(true)` makes this a true @@ -582,7 +570,7 @@ impl LsmVectorSearchPlanner { // generation resolves its own schema before scanning, and // the inlined future is too deep for the `Send` proof. let reconciled = generation.reconcile(Box::pin(scanner.create_plan()).await?)?; - match above { + match &above { Some(expr) => filter_above(reconciled, expr), None => Ok(reconciled), } diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index be9592b35fd..8167f55aa06 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -1602,29 +1602,6 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// Re-label `batch` to the storage schema, matching columns by **field id** -/// where both sides carry one, and by **name** otherwise. -/// -/// A column the schema declares and the batch does not carry is filled with -/// typed nulls; `_tombstone` is filled with `false`. A column the batch carries -/// and the schema does not declare is dropped. -/// -/// Ids are tried first because a rename keeps the id and changes the name: a -/// name match would null the new name and drop the old one, losing the column's -/// values. An entry carrying no ids falls back to the name match. -/// -/// Both cases are what a replayed WAL entry looks like after the table's schema -/// moved: an entry predates a column added since, and carries one dropped -/// since. Live writes reach here already checked against the logical schema, so -/// for them every column is present and this only appends `_tombstone`. -/// -/// A column whose type changed is cast, the same way `alter_columns` casts the -/// base data, so a replayed row lands in the state it would have had if it had -/// been written after the change. A cast that would lose information is an -/// error, not a silent null. -/// -/// A primary key the batch does not carry stays an error — there is no value to -/// invent. /// A batch a caller just handed in, under the storage schema. /// /// Live input is trusted for its values and not for its identity: it has been @@ -1650,6 +1627,27 @@ fn conform_live_batch( conform_to_storage_schema(batch, storage_schema, pk_columns) } +/// Re-label `batch` to the storage schema, matching columns by **field id** +/// where both sides carry one, and by **name** otherwise. +/// +/// A column the schema declares and the batch does not carry is filled with +/// typed nulls; `_tombstone` is filled with `false`. A column the batch carries +/// and the schema does not declare is dropped. +/// +/// Ids are tried first: a name match would null the new name and drop the old +/// one, losing the column's values. An entry carrying no ids falls back to the +/// name match. +/// +/// Both are what a replayed WAL entry looks like after the table's schema +/// moved: it predates a column added since, and carries one dropped since. A +/// live write arrives already checked against the logical schema, so for it +/// every column is present and this only appends `_tombstone`. +/// +/// A column whose scalar type has moved is an error rather than a cast: a table +/// with a MemWAL refuses a retype, so a disagreement here is one to surface. +/// +/// A primary key the batch does not carry is an error — there is no value to +/// invent. fn conform_to_storage_schema( batch: RecordBatch, storage_schema: &Arc, @@ -2137,10 +2135,9 @@ impl ShardWriter { // The caller's schema is the shard's logical schema; the storage schema // is derived below, once the primary key is known. lance owns // `_tombstone` and appends it here — idempotent across reopens. - // The stored schema carries field ids so identity survives a rename; - // what a caller's batch is checked against must not, since a batch - // carries none and Arrow compares a struct's children in full. let tombstoned = schema_with_tombstone(&schema); + // What a caller's batch is checked against carries no ids: a batch has + // none, and Arrow compares a struct's children in full. let logical_schema = Arc::new(without_field_ids(&schema)); let base_uri = base_uri.into(); @@ -4877,11 +4874,12 @@ mod tests { ); } - /// A struct column is matched whole, by the id on the column itself. + /// A struct column renamed at the top level is matched by the id on the + /// column itself, and taken whole. /// - /// Ids are carried for top-level fields, which is the granularity conform - /// works at: a column is taken or it is not. A change inside the struct - /// changes the column's type, and is handled as a type change. + /// Only the parent moves here. Reconciliation is recursive — a child + /// carries its own id and is resolved on its own — which the nested cases + /// in `reconcile` cover. #[test] fn test_conform_matches_a_struct_column_by_its_own_id() { fn with_id(field: ArrowField, id: i32) -> ArrowField { diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 58ae3fee226..f056aec4553 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -472,32 +472,31 @@ pub(super) async fn add_columns( .await } -/// Refuse to change a column's type on a table with a MemWAL attached. -/// -/// A cast gives the column a new field id and keeps its name, which is exactly -/// what dropping a column and adding another under that name looks like. Rows -/// the WAL still holds carry the old id, and nothing in the schemas says which -/// of the two happened -- so they could only be reconciled by guessing. -/// -/// A table without a MemWAL is unaffected: this is the only thing the check -/// looks at. -/// -/// Takes the decision as a bool rather than the alterations themselves: a -/// reference to them held across an await would have to be `Sync`. /// What `alter_columns` refuses on a table with a MemWAL, and why. /// -/// Both are alterations whose meaning depends on rows the commit cannot see. +/// Each is an alteration whose meaning depends on rows the commit cannot see. enum Unsupported { /// A cast takes a new field id, which rows still in the WAL cannot be /// matched to. Retype, - /// Renaming a nested child onto a name one of its siblings currently holds. - /// Compaction matches a struct's children by name, so both sides would - /// carry the same two names and the merge could not tell which child is - /// which -- it would write each one's values under the other's name. - /// Resolving nested children by id there is more machinery than this case - /// is worth; a rename that does not collide with a sibling is unaffected. - SiblingNameReuse, + /// Renaming a field inside a struct. + /// + /// Compaction relabels a generation's columns by field id at the top level + /// only, so a struct's children are merged under the names the generation + /// stored. A generation sealed before the rename still carries the old + /// child name, and the merge has no id to follow. + /// + /// One rename makes the two shapes disagree, which the merge refuses. Two + /// that exchange a pair of names -- reachable in three commits through a + /// free name -- leave the shapes identical and the identities crossed, and + /// the merge writes each child's values under the other's name with the API + /// reporting success throughout. + /// + /// Refused as a whole rather than only on a collision: a collision check + /// sees the schema as it stands and cannot see the history that reaches the + /// same place. Re-enabling this needs compaction to resolve a struct's + /// children by id. + NestedRename, /// Tightening a column to non-null is validated against the committed /// fragments, and a row still in the MemWAL is not among them. Flushing /// first does not close the window: a flush covers the generations open @@ -507,30 +506,30 @@ enum Unsupported { Tightening, } -/// Whether `alteration` renames a nested field onto a name one of its siblings -/// currently holds. See [`Unsupported::SiblingNameReuse`]. +/// Whether `alteration` renames a field inside a struct. +/// See [`Unsupported::NestedRename`]. /// /// Top-level renames are exempt: compaction resolves those by field id, so a /// pair of them can exchange names safely. -fn renames_onto_a_sibling(dataset: &Dataset, alteration: &ColumnAlteration) -> bool { - let Some(rename) = &alteration.rename else { +fn renames_a_nested_field(dataset: &Dataset, alteration: &ColumnAlteration) -> bool { + if alteration.rename.is_none() { return false; - }; + } // Resolved rather than split on `.`: a field name may contain dots, in - // which case the path quotes it, and splitting would name the wrong parent - // and admit exactly the rename this refuses. + // which case the path quotes it, and a split would misjudge the depth. let Some(chain) = dataset.schema().resolve(&alteration.path) else { return false; }; - let [.., parent, child] = chain.as_slice() else { - return false; - }; - parent - .children - .iter() - .any(|sibling| sibling.name != child.name && sibling.name == *rename) + chain.len() > 1 } +/// Refuse `unsupported` when the table has a MemWAL attached. +/// +/// A table without one is unaffected: the presence of a MemWAL is the only +/// thing this looks at. +/// +/// Takes the decision already made rather than the alterations themselves: a +/// reference to them held across the await would have to be `Sync`. async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) -> Result<()> { let Some(unsupported) = unsupported else { return Ok(()); @@ -549,11 +548,10 @@ async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) runs against the base table, and a write admitted into the WAL while it runs is \ not there to be checked. Drop the MemWAL first." } - Unsupported::SiblingNameReuse => { - "cannot rename a nested field onto a name one of its siblings holds on a table \ - with a MemWAL attached: rows still in the WAL would have the two children \ - matched by name and their values exchanged. Rename the sibling out of the way \ - first, or drop the MemWAL." + Unsupported::NestedRename => { + "cannot rename a field inside a struct on a table with a MemWAL attached: \ + compaction matches a struct's children by name, and rows still in the WAL \ + carry the old one. Drop the MemWAL first." } })) } @@ -822,9 +820,9 @@ pub(super) async fn alter_columns( Some(Unsupported::Tightening) } else if alterations .iter() - .any(|a| renames_onto_a_sibling(dataset, a)) + .any(|a| renames_a_nested_field(dataset, a)) { - Some(Unsupported::SiblingNameReuse) + Some(Unsupported::NestedRename) } else { None }; @@ -4635,6 +4633,61 @@ mod test { Ok(()) } + /// A swap reached through a free name is refused at its first step. + /// + /// Each step looks harmless against the schema as it stands, so a collision + /// check admits all three and the pair ends up exchanged -- while a + /// generation sealed beforehand still has them the other way round, and + /// compaction would merge each child's values under the other's name. + #[tokio::test] + async fn a_swap_through_a_free_name_is_refused_with_a_mem_wal() -> Result<()> { + let a = ArrowField::new("a", DataType::Int32, true); + let b = ArrowField::new("b", DataType::Int32, true); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "parent", + DataType::Struct(ArrowFields::from(vec![a.clone(), b.clone()])), + true, + )])); + let parent = StructArray::new( + ArrowFields::from(vec![a, b]), + vec![ + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + Arc::new(Int32Array::from(vec![2])) as ArrayRef, + ], + None, + ); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(parent)])?; + let test_dir = TempStrDir::default(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, &test_dir, None).await?; + dataset.initialize_mem_wal().execute().await?; + + // The first step of the detour targets a free name and is still refused. + let err = dataset + .alter_columns(&[ColumnAlteration::new("parent.a".into()).rename("tmp".into())]) + .await + .expect_err("a nested rename must be refused while a MemWAL is attached"); + assert!( + err.to_string().contains("rename a field inside a struct"), + "unexpected error: {err}" + ); + + // The schema is untouched, so no later step can reach the swap. + let DataType::Struct(children) = dataset.schema().field("parent").unwrap().data_type() + else { + panic!("parent is a struct"); + }; + let names: Vec<&str> = children.iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, vec!["a", "b"], "the refusal left the pair as it was"); + + // A top-level rename is unaffected: compaction resolves those by id. + dataset + .alter_columns(&[ColumnAlteration::new("parent".into()).rename("outer".into())]) + .await?; + assert!(dataset.schema().field("outer").is_some()); + Ok(()) + } + /// A field name may contain dots, in which case the path quotes it. /// Splitting on the last dot names the wrong parent, and the sibling check /// then admits the rename it exists to refuse. @@ -4663,15 +4716,14 @@ mod test { let onto_sibling = ColumnAlteration::new("parent.`child.with.dot`".into()).rename("sibling".into()); assert!( - renames_onto_a_sibling(&dataset, &onto_sibling), - "renaming onto a sibling's name must be seen" + renames_a_nested_field(&dataset, &onto_sibling), + "a rename inside a struct must be seen through a quoted path" ); - let onto_free = - ColumnAlteration::new("parent.`child.with.dot`".into()).rename("free".into()); + let top_level = ColumnAlteration::new("parent".into()).rename("free".into()); assert!( - !renames_onto_a_sibling(&dataset, &onto_free), - "a name no sibling holds is not a collision" + !renames_a_nested_field(&dataset, &top_level), + "a top-level rename is not a nested one" ); Ok(()) } From d6fb92173b068be3b317ad7a5f48126df23d28fb Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 15:04:13 -0400 Subject: [PATCH 45/48] feat(mem_wal): offer the generation resolution to a caller that compacts A sealed generation's columns relate to the table's by field id, and a struct's children carry ids of their own. A caller merging a generation into the base table by matching names cannot follow a rename: one leaves the two shapes disagreeing, and a pair that exchange names leaves the shapes identical and the values crossed. `reconcile_batches` hands that caller the resolution the read paths already apply, so the merge follows ids at every level. With the resolution available, refusing a nested rename is no longer the only way to keep one safe, so the refusal is withdrawn. A retype and a nullability tightening stay refused: those are undecidable from schemas alone whatever the caller does. --- rust/lance/src/dataset/mem_wal.rs | 107 ++++++++++++++++ rust/lance/src/dataset/schema_evolution.rs | 140 --------------------- 2 files changed, 107 insertions(+), 140 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index af7df745dc5..b854a22458b 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -51,6 +51,7 @@ use std::sync::Arc; use lance_core::datatypes::{Field, LANCE_FIELD_ID_KEY, Schema}; +use arrow_array::RecordBatch; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; /// Column name for the mem_wal tombstone (delete sentinel) marker. @@ -207,6 +208,36 @@ pub fn schema_with_tombstone(base: &ArrowSchema) -> Arc { )) } +/// `batches`, written under `source_schema`, brought to `target_schema`. +/// +/// Columns are matched by field id where both sides carry one, and by name +/// otherwise. A column `target_schema` declares and the batches do not carry is +/// filled with typed nulls; `_tombstone` is filled with `false`. A column the +/// batches carry and `target_schema` does not declare is dropped. A primary key +/// the batches do not carry is an error, since no value can stand in for it. +/// +/// This is the same resolution a read of a sealed generation applies, offered +/// to a caller that reads one for itself. Matching nested children by name +/// cannot follow a rename: a struct's children carry ids of their own, and only +/// those relate a generation's copy of a column to the table's. +/// +/// `batches` must be in `source_schema`'s column order, as a scan of the +/// dataset it describes returns them. The result is in `target_schema`'s order, +/// under a schema carrying no field ids. +pub fn reconcile_batches( + source_schema: &ArrowSchema, + target_schema: &Arc, + pk_columns: &[String], + batches: Vec, +) -> lance_core::Result> { + let plan = + reconcile::Plan::resolve(source_schema, target_schema, pk_columns)?.emitting_plain_schema(); + if plan.is_identity() { + return Ok(batches); + } + batches.iter().map(|batch| plan.apply(batch)).collect() +} + pub use api::{DatasetMemWalExt, InitializeMemWalBuilder, validate_maintained_indexes}; pub use index::{MemIndexKind, MemTableVisibility}; pub use manifest::ShardManifestStore; @@ -235,6 +266,82 @@ mod tests { ]) } + fn stamped(name: &str, data_type: DataType, id: i32) -> ArrowField { + ArrowField::new(name, data_type, true).with_metadata( + [(LANCE_FIELD_ID_KEY.to_string(), id.to_string())] + .into_iter() + .collect(), + ) + } + + /// Two children exchanging names is the case a name match cannot survive: + /// both sides carry the same two names, so only the ids say which values + /// belong to which. Each child's values must follow its id to the name the + /// target now gives it. + #[test] + fn a_pair_of_children_that_swapped_names_follow_their_ids() { + let struct_of = |first: &str, second: &str, ids: (i32, i32)| { + DataType::Struct(Fields::from(vec![ + stamped(first, DataType::Int64, ids.0), + stamped(second, DataType::Int64, ids.1), + ])) + }; + let source = ArrowSchema::new(vec![ + stamped("id", DataType::Int64, 0), + stamped("info", struct_of("a", "b", (1, 2)), 3), + ]); + // The table has since exchanged the two children's names; the ids stay. + let target = Arc::new(ArrowSchema::new(vec![ + stamped("id", DataType::Int64, 0), + stamped("info", struct_of("b", "a", (1, 2)), 3), + ])); + + let info = arrow_array::StructArray::new( + match source.field(1).data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!("info is a struct"), + }, + vec![ + Arc::new(arrow_array::Int64Array::from(vec![10])) as arrow_array::ArrayRef, + Arc::new(arrow_array::Int64Array::from(vec![20])), + ], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(source.clone()), + vec![ + Arc::new(arrow_array::Int64Array::from(vec![1])), + Arc::new(info), + ], + ) + .expect("a batch under the source schema"); + + let out = reconcile_batches(&source, &target, &["id".to_string()], vec![batch]) + .expect("reconcile"); + let info = out[0] + .column(1) + .as_any() + .downcast_ref::() + .expect("info is a struct"); + + // `b` is the name id 1 now wears, so it must hold id 1's value. + let b = info + .column_by_name("b") + .expect("b") + .as_any() + .downcast_ref::() + .expect("int64"); + assert_eq!(b.value(0), 10, "id 1's value follows its id to `b`"); + + let a = info + .column_by_name("a") + .expect("a") + .as_any() + .downcast_ref::() + .expect("int64"); + assert_eq!(a.value(0), 20, "id 2's value follows its id to `a`"); + } + #[test] fn relax_widens_every_non_pk_field_and_leaves_the_key_alone() { let relaxed = relax_non_pk_nullability(&logical(), &["id".to_string()]); diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index f056aec4553..07bb80d780e 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -479,24 +479,6 @@ enum Unsupported { /// A cast takes a new field id, which rows still in the WAL cannot be /// matched to. Retype, - /// Renaming a field inside a struct. - /// - /// Compaction relabels a generation's columns by field id at the top level - /// only, so a struct's children are merged under the names the generation - /// stored. A generation sealed before the rename still carries the old - /// child name, and the merge has no id to follow. - /// - /// One rename makes the two shapes disagree, which the merge refuses. Two - /// that exchange a pair of names -- reachable in three commits through a - /// free name -- leave the shapes identical and the identities crossed, and - /// the merge writes each child's values under the other's name with the API - /// reporting success throughout. - /// - /// Refused as a whole rather than only on a collision: a collision check - /// sees the schema as it stands and cannot see the history that reaches the - /// same place. Re-enabling this needs compaction to resolve a struct's - /// children by id. - NestedRename, /// Tightening a column to non-null is validated against the committed /// fragments, and a row still in the MemWAL is not among them. Flushing /// first does not close the window: a flush covers the generations open @@ -506,23 +488,6 @@ enum Unsupported { Tightening, } -/// Whether `alteration` renames a field inside a struct. -/// See [`Unsupported::NestedRename`]. -/// -/// Top-level renames are exempt: compaction resolves those by field id, so a -/// pair of them can exchange names safely. -fn renames_a_nested_field(dataset: &Dataset, alteration: &ColumnAlteration) -> bool { - if alteration.rename.is_none() { - return false; - } - // Resolved rather than split on `.`: a field name may contain dots, in - // which case the path quotes it, and a split would misjudge the depth. - let Some(chain) = dataset.schema().resolve(&alteration.path) else { - return false; - }; - chain.len() > 1 -} - /// Refuse `unsupported` when the table has a MemWAL attached. /// /// A table without one is unaffected: the presence of a MemWAL is the only @@ -548,11 +513,6 @@ async fn reject_on_mem_wal(dataset: &Dataset, unsupported: Option) runs against the base table, and a write admitted into the WAL while it runs is \ not there to be checked. Drop the MemWAL first." } - Unsupported::NestedRename => { - "cannot rename a field inside a struct on a table with a MemWAL attached: \ - compaction matches a struct's children by name, and rows still in the WAL \ - carry the old one. Drop the MemWAL first." - } })) } @@ -818,11 +778,6 @@ pub(super) async fn alter_columns( Some(Unsupported::Retype) } else if alterations.iter().any(|a| a.nullable == Some(false)) { Some(Unsupported::Tightening) - } else if alterations - .iter() - .any(|a| renames_a_nested_field(dataset, a)) - { - Some(Unsupported::NestedRename) } else { None }; @@ -4632,99 +4587,4 @@ mod test { Ok(()) } - - /// A swap reached through a free name is refused at its first step. - /// - /// Each step looks harmless against the schema as it stands, so a collision - /// check admits all three and the pair ends up exchanged -- while a - /// generation sealed beforehand still has them the other way round, and - /// compaction would merge each child's values under the other's name. - #[tokio::test] - async fn a_swap_through_a_free_name_is_refused_with_a_mem_wal() -> Result<()> { - let a = ArrowField::new("a", DataType::Int32, true); - let b = ArrowField::new("b", DataType::Int32, true); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "parent", - DataType::Struct(ArrowFields::from(vec![a.clone(), b.clone()])), - true, - )])); - let parent = StructArray::new( - ArrowFields::from(vec![a, b]), - vec![ - Arc::new(Int32Array::from(vec![1])) as ArrayRef, - Arc::new(Int32Array::from(vec![2])) as ArrayRef, - ], - None, - ); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(parent)])?; - let test_dir = TempStrDir::default(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); - let mut dataset = Dataset::write(reader, &test_dir, None).await?; - dataset.initialize_mem_wal().execute().await?; - - // The first step of the detour targets a free name and is still refused. - let err = dataset - .alter_columns(&[ColumnAlteration::new("parent.a".into()).rename("tmp".into())]) - .await - .expect_err("a nested rename must be refused while a MemWAL is attached"); - assert!( - err.to_string().contains("rename a field inside a struct"), - "unexpected error: {err}" - ); - - // The schema is untouched, so no later step can reach the swap. - let DataType::Struct(children) = dataset.schema().field("parent").unwrap().data_type() - else { - panic!("parent is a struct"); - }; - let names: Vec<&str> = children.iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, vec!["a", "b"], "the refusal left the pair as it was"); - - // A top-level rename is unaffected: compaction resolves those by id. - dataset - .alter_columns(&[ColumnAlteration::new("parent".into()).rename("outer".into())]) - .await?; - assert!(dataset.schema().field("outer").is_some()); - Ok(()) - } - - /// A field name may contain dots, in which case the path quotes it. - /// Splitting on the last dot names the wrong parent, and the sibling check - /// then admits the rename it exists to refuse. - #[tokio::test] - async fn a_quoted_path_resolves_to_its_real_parent() -> Result<()> { - let child = ArrowField::new("child.with.dot", DataType::Int32, true); - let sibling = ArrowField::new("sibling", DataType::Int32, true); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "parent", - DataType::Struct(ArrowFields::from(vec![child.clone(), sibling.clone()])), - true, - )])); - let parent = StructArray::new( - ArrowFields::from(vec![child, sibling]), - vec![ - Arc::new(Int32Array::from(vec![1])) as ArrayRef, - Arc::new(Int32Array::from(vec![2])) as ArrayRef, - ], - None, - ); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(parent)])?; - let test_dir = TempStrDir::default(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); - let dataset = Dataset::write(reader, &test_dir, None).await?; - - let onto_sibling = - ColumnAlteration::new("parent.`child.with.dot`".into()).rename("sibling".into()); - assert!( - renames_a_nested_field(&dataset, &onto_sibling), - "a rename inside a struct must be seen through a quoted path" - ); - - let top_level = ColumnAlteration::new("parent".into()).rename("free".into()); - assert!( - !renames_a_nested_field(&dataset, &top_level), - "a top-level rename is not a nested one" - ); - Ok(()) - } } From 29423eb8fdf3630fa7e99ae72c87df234b58f5fa Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 15:42:42 -0400 Subject: [PATCH 46/48] fix(mem_wal): keep a generation's own ids out of the resolution A generation numbers `_tombstone` in its own schema, so its id is whatever that generation reached -- and the table may have given that same number to a column of its own. Honouring it resolves the two to each other and refuses the merge on their types. The scan path already strips those before resolving; `reconcile_batches` now does the same, so a caller reading a generation for itself gets the same answer a read of it would give. --- rust/lance/src/dataset/mem_wal.rs | 60 ++++++++++++++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index b854a22458b..527a432f206 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -230,8 +230,25 @@ pub fn reconcile_batches( pk_columns: &[String], batches: Vec, ) -> lance_core::Result> { + // A generation numbers its own columns -- `_tombstone`, and anything the + // table has since dropped -- in its own schema, so those ids collide with + // whatever the table gave those numbers. Stripped before resolution, or a + // column added to the table resolves to whichever of them shares its id. + let source = ArrowSchema::new_with_metadata( + source_schema + .fields() + .iter() + .map(|field| { + match field.name() != TOMBSTONE && !lance_core::is_system_column(field.name()) { + true => field.as_ref().clone(), + false => reconcile::without_field_id(field), + } + }) + .collect::>(), + source_schema.metadata().clone(), + ); let plan = - reconcile::Plan::resolve(source_schema, target_schema, pk_columns)?.emitting_plain_schema(); + reconcile::Plan::resolve(&source, target_schema, pk_columns)?.emitting_plain_schema(); if plan.is_identity() { return Ok(batches); } @@ -274,6 +291,47 @@ mod tests { ) } + /// A generation numbers `_tombstone` in its own schema, so its id is + /// whatever that generation reached -- and the table has given that same + /// number to a column of its own. Honouring it would resolve the two to + /// each other and refuse the merge on their types. + #[test] + fn a_generations_tombstone_does_not_answer_for_a_column_sharing_its_id() { + let source = ArrowSchema::new(vec![ + stamped("id", DataType::Int64, 0), + stamped(TOMBSTONE, DataType::Boolean, 1), + ]); + // The table gave id 1 to a column added after that generation sealed. + let target = Arc::new(ArrowSchema::new(vec![ + stamped("id", DataType::Int64, 0), + stamped("extra", DataType::Int64, 1), + ArrowField::new(TOMBSTONE, DataType::Boolean, true), + ])); + let batch = RecordBatch::try_new( + Arc::new(source.clone()), + vec![ + Arc::new(arrow_array::Int64Array::from(vec![1])), + Arc::new(arrow_array::BooleanArray::from(vec![false])), + ], + ) + .expect("a batch under the source schema"); + + let out = reconcile_batches(&source, &target, &["id".to_string()], vec![batch]) + .expect("the tombstone's id must not be honoured"); + let out = &out[0]; + assert!( + out.column_by_name("extra").expect("extra").is_null(0), + "the added column has no value in a generation sealed before it" + ); + let tombstone = out + .column_by_name(TOMBSTONE) + .expect("_tombstone") + .as_any() + .downcast_ref::() + .expect("boolean"); + assert!(!tombstone.value(0), "and the row is still live"); + } + /// Two children exchanging names is the case a name match cannot survive: /// both sides carry the same two names, so only the ids say which values /// belong to which. Each child's values must follow its id to the name the From 0f4e0790149564f1a4dd4589e0a184f84d02c68a Mon Sep 17 00:00:00 2001 From: XYZhan Date: Wed, 16 Sep 2026 17:40:14 -0400 Subject: [PATCH 47/48] test(mem_wal): fold the struct push-down cases into one table Three tests built the same generation two ways and asserted whether a predicate on the parent reached it. They are one rule with three inputs, so they read better as a table: unmoved pushes down, a renamed or a replaced child does not. --- .../src/dataset/mem_wal/scanner/generation.rs | 103 ++++++++---------- 1 file changed, 43 insertions(+), 60 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index d0cf1aa045e..395e3b9a95e 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -426,6 +426,49 @@ mod tests { ) } + /// A predicate naming a struct can only be pushed down when the generation + /// stores that struct exactly as the table declares it. A nested reference + /// names the parent, and a parent's name does not move when one of its + /// children does, so the parent's name alone cannot say whether pushing + /// down is safe. + #[test] + fn a_struct_predicate_is_pushed_down_only_when_its_children_did_not_move() { + // The child as the generation stored it, as the table now declares it, + // and whether a predicate on the parent may reach the stored data. + let cases = [ + // Deferring this one was the bug: a search takes its top-k first, + // so a predicate applied afterwards loses a lower-ranked row that + // should have won. + ("nothing moved", ("c", 2), ("c", 2), true), + ("the child was renamed", ("c", 2), ("d", 2), false), + // Same name, same type, different field id: a child dropped and + // added back is a different column wearing the old one's shape, so + // a predicate on it must not reach the retired values. + ("the child was replaced", ("c", 2), ("c", 9), false), + ]; + for (case, stored_child, table_child, pushed_down) in cases { + let info = |(name, id): (&str, i32)| { + with_id( + "info", + DataType::Struct(Fields::from(vec![with_id(name, DataType::Int64, id)])), + 1, + ) + }; + let read = generation( + schema(vec![with_id("id", DataType::Int64, 0), info(stored_child)]), + schema(vec![with_id("id", DataType::Int64, 0), info(table_child)]), + &["id", "info"], + ); + let expr = col("info").is_not_null(); + let got = read.to_stored(&expr); + assert_eq!( + got, + pushed_down.then_some(expr), + "{case}: the predicate was pushed down when it should not have been, or the reverse" + ); + } + } + #[test] fn a_renamed_column_is_asked_for_under_the_name_the_generation_has() { assert_eq!(renamed().stored_projection(), vec!["id", "value"]); @@ -501,66 +544,6 @@ mod tests { assert_eq!(read.to_stored(&col("added").eq(lit(1i64))), None); } - /// A nested reference names the parent, and a parent's name does not move - /// when a child is renamed — so pushing it down would evaluate it against - /// child names the table does not have. - #[test] - fn a_predicate_on_a_struct_whose_child_was_renamed_is_not_pushed_down() { - let nested = |child: &str| { - with_id( - "info", - DataType::Struct(Fields::from(vec![with_id(child, DataType::Int64, 2)])), - 1, - ) - }; - let read = generation( - schema(vec![with_id("id", DataType::Int64, 0), nested("c")]), - schema(vec![with_id("id", DataType::Int64, 0), nested("d")]), - &["id", "info"], - ); - assert_eq!(read.to_stored(&col("info").is_not_null()), None); - } - - /// The common case: a struct nothing moved is pushed down like any other - /// column. Deferring these was the bug -- a search takes its top-k first, - /// so a predicate applied afterwards loses a lower-ranked row that should - /// have won. - #[test] - fn a_predicate_on_a_struct_that_did_not_move_is_pushed_down() { - let unmoved = with_id( - "info", - DataType::Struct(Fields::from(vec![with_id("c", DataType::Int64, 2)])), - 1, - ); - let read = generation( - schema(vec![with_id("id", DataType::Int64, 0), unmoved.clone()]), - schema(vec![with_id("id", DataType::Int64, 0), unmoved]), - &["id", "info"], - ); - let expr = col("info").is_not_null(); - assert_eq!(read.to_stored(&expr), Some(expr)); - } - - /// Same name, same type, different field id: a child dropped and added - /// back is a different column wearing the old one's shape, so a predicate - /// on it must not reach the retired values. - #[test] - fn a_predicate_on_a_struct_whose_child_was_replaced_is_not_pushed_down() { - let child = |id: i32| { - with_id( - "info", - DataType::Struct(Fields::from(vec![with_id("c", DataType::Int64, id)])), - 1, - ) - }; - let read = generation( - schema(vec![with_id("id", DataType::Int64, 0), child(2)]), - schema(vec![with_id("id", DataType::Int64, 0), child(9)]), - &["id", "info"], - ); - assert_eq!(read.to_stored(&col("info").is_not_null()), None); - } - /// With no ids to match on, the table's own names are the only link — the /// behaviour a caller that supplies no identity schema gets. #[test] From 4be705162874c7a5f7797c68b443b44217631973 Mon Sep 17 00:00:00 2001 From: XYZhan Date: Thu, 17 Sep 2026 02:21:12 -0400 Subject: [PATCH 48/48] docs(mem_wal): state the field-id invariant the resolution rests on `max_field_id` is the maximum over the current schema and the fields the base fragments reference, so dropping a column can lower it and let the next added column take the id back -- while a generation retained across that sequence still holds the dropped column under it. Resolving by id then reads the retired values as the column that took the id, and nothing in either schema tells that apart from a rename. Enforcing it belongs where ids are allocated, which is not this module. Documented here with a test pinning what happens when it is not kept, so a consumer knows the discipline it owes and a fix has something to flip. --- .../src/dataset/mem_wal/scanner/generation.rs | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/scanner/generation.rs b/rust/lance/src/dataset/mem_wal/scanner/generation.rs index 395e3b9a95e..4e93171f12b 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/generation.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -36,6 +36,21 @@ use crate::dataset::mem_wal::{TOMBSTONE, arrow_schema_with_field_ids}; /// pushed into it and under which names ([`Self::to_stored`]), and how to /// bring the result back to the table's names and shapes /// ([`Self::reconcile`]). +/// +/// # The invariant this rests on +/// +/// A field id identifies a column for as long as any generation holds it. +/// Lance does not enforce that: `max_field_id` is the maximum over the current +/// schema and the fields the base fragments reference, so dropping a column can +/// lower it and let the next added column take the id back -- while a retained +/// generation still holds the dropped column under it. Resolving by id then +/// reads the retired values as the column that took the id, and nothing in +/// either schema tells the two apart from a rename. +/// +/// Keeping it is therefore the caller's: a consumer that retains generations +/// must drain them before a drop/add sequence can reuse an id they still hold. +/// `a_field_id_taken_back_after_a_drop_is_read_as_the_column_that_took_it` +/// pins what happens when it is not kept. pub(super) struct GenerationRead { /// The generation's own schema, carrying its field ids. stored_schema: Schema, @@ -475,6 +490,46 @@ mod tests { assert_eq!(renamed().stored_name("amount"), Some("value")); } + /// A KNOWN LIMITATION, pinned so a fix is visible when it lands. + /// + /// Lance's `max_field_id` is not a permanent high-water mark: it is the + /// maximum over the current schema and the fields the base fragments + /// reference, so dropping a column can lower it and let the next added + /// column take the id back. A generation retained across that sequence + /// still holds the dropped column under that id, and this resolution -- + /// which follows ids, by design -- then reads it as the column that took + /// the id. Nothing in either schema distinguishes that from a rename, and + /// when the types agree the values are returned under the new name rather + /// than as null. `reconcile_batches` makes the same pairing, so a merge + /// persists them. + /// + /// The invariant this rests on is therefore the caller's to keep: a + /// consumer that retains generations must not let a field id be reused + /// while one still holds it -- draining the retained generations before a + /// drop/add sequence is what establishes that. + #[test] + fn a_field_id_taken_back_after_a_drop_is_read_as_the_column_that_took_it() { + let read = generation( + // Sealed while id 1 was `retired`. + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("retired", DataType::Int64, 1), + ]), + // `retired` was dropped and `added` took its id back. + schema(vec![ + with_id("id", DataType::Int64, 0), + with_id("added", DataType::Int64, 1), + ]), + &["id", "added"], + ); + assert_eq!( + read.stored_name("added"), + Some("retired"), + "today the reused id pairs them; a fix makes this `None`, and the \ + generation contributes null for `added` instead" + ); + } + #[test] fn a_column_the_generation_never_stored_is_left_out_of_the_projection() { let read = generation(