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..527a432f206 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,6 +49,9 @@ pub mod write; 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. @@ -113,6 +117,79 @@ pub fn relax_non_pk_nullability( )) } +/// 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 fn arrow_schema_with_field_ids(schema: &Schema) -> ArrowSchema { + let arrow: ArrowSchema = schema.into(); + let fields: Vec = arrow + .fields() + .iter() + .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.clone().with_metadata(metadata) + } + _ => field.clone(), + }; + // 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)) + } + 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, + } +} + /// Extend the logical schema with the trailing `_tombstone` column — the /// intermediate [`relax_non_pk_nullability`] widens into the storage schema. /// @@ -131,6 +208,53 @@ 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> { + // 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, 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; @@ -159,6 +283,123 @@ 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(), + ) + } + + /// 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 + /// 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/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index ef55d20d5ab..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; @@ -678,8 +683,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; @@ -702,13 +712,26 @@ 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 } } +/// 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, +} + /// Build the in-memory index configurations for `index_names`. /// /// Shared by [`DatasetMemWalExt::mem_wal_writer`] and @@ -718,6 +741,7 @@ 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 +753,26 @@ 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 does not have: + // 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!( "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. 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..c3914acc9de --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/reconcile.rs @@ -0,0 +1,869 @@ +// 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 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; + +use arrow::array::ArrayData; +use arrow_array::{ + 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; +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() + .get(LANCE_FIELD_ID_KEY) + .and_then(|v| v.parse::().ok()) + .filter(|id| *id >= 0) +} + +/// `column` under `data_type`, which differs from its own only in the field +/// ids it carries. +/// +/// 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()); + } + relabel_data(&column.to_data(), data_type).map(arrow_array::make_array) +} + +/// [`relabel_to`] over one array's data, and its children's in turn. Arrow +/// validates a container against the child types its own type declares, so a +/// label that moved at any depth has to move at every level below it. +fn relabel_data(data: &ArrayData, data_type: &DataType) -> 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}", + data.data_type() + )) + }) +} + +/// [`without_field_ids`] for a type rather than a schema, for the nested types a +/// reconciliation builds. +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() +} + +/// `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 field = without_field_id(field); + 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())) + } + 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)) + } + DataType::Map(entries, sorted) => { + let sorted = *sorted; + field + .clone() + .with_data_type(DataType::Map(Arc::new(strip(entries)), sorted)) + } + _ => 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 { + /// The same plan emitting the table's plain Arrow schema. + /// + /// 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 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 { + 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 + /// 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.fields(), target.fields()); + 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}"))) + } +} + +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)) + }); + // 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() + .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]; + // 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)); + } + // 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() + ))); + } + 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 { + children_of(data_type).is_some() +} + +/// 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()), + // 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, + } +} + +/// 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_by_id(&source_children, &target_children); + target_children + .iter() + .map(|child| resolve_field(child, &source_children, &claimed, &[])) + .collect() +} + +/// 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() + .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 +} + +/// 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}")))?, + )) +} + +/// 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])), + // 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 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 { + 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 or list 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]))), + } +} + +#[cfg(test)] +mod relabel_tests { + use super::*; + use arrow_array::{ + Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, StructArray, + }; + 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())] + .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 + } + + /// 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); + 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); + } + + /// 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] + 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); + } +} 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 dbd5e9e4a79..b78615f9512 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,9 @@ impl LsmScanner { pk_columns: Vec, ) -> Self { Self { + // 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, shard_snapshots, @@ -339,6 +348,18 @@ impl LsmScanner { self } + /// 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). + /// + /// 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 + } + /// 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 @@ -585,6 +606,7 @@ impl LsmScanner { 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()); @@ -652,6 +674,7 @@ impl LsmScanner { let collector = self.build_collector(); 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()); @@ -697,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()); } @@ -719,7 +743,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..46de19e87d3 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(crate) 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/exec/schema_relabel.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs index 88691f99d63..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}; @@ -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 8092d1b0329..cc6292f6d9a 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -58,6 +58,7 @@ use super::block_list::compute_source_block_lists; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; use super::exec::{FirstByPkExec, PkBlockFilterExec}; +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::memtable::scanner::MemTableScanner; @@ -475,6 +476,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. @@ -501,6 +505,7 @@ impl LsmFtsSearchPlanner { Self { collector, pk_columns, + identity_schema: base_schema.clone(), base_schema, session: None, store_params: None, @@ -527,6 +532,16 @@ impl LsmFtsSearchPlanner { self } + /// 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 + } + /// Set the session used to open SSTables. pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); @@ -832,16 +847,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, - columns, - &query, - *fetch_limit, - projection, - &index_params, - )) + let arm: futures::future::BoxFuture<'_, Result>> = + Box::pin(self.build_source_plan( + source, + columns, + &query, + *fetch_limit, + projection, + &index_params, + &target_schema, + )); + arm })) .await?; @@ -926,14 +945,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(), + &self.identity_schema, + &self.pk_columns, + 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) @@ -1003,6 +1038,9 @@ impl LsmFtsSearchPlanner { limit: Option, projection: Option<&[String]>, index_params: &HashMap<&str, InvertedIndexParams>, + // What every arm is normalized to, so an arm with nothing to offer can + // stand in for itself. + target_schema: &SchemaRef, ) -> Result> { // One column: bind every leaf to it, which is what lets a tree built // from bare terms reach the right field. Several: the leaves already @@ -1047,16 +1085,87 @@ 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::>())?; - if let Some(ref filter) = self.filter { + // Asked of this generation under its own names: a rename moved + // the table's name while the file still holds the old one. + let asked_for = self.fts_scanner_projection(projection); + let mut generation = GenerationRead::new( + dataset.schema(), + &self.identity_schema, + &self.pk_columns, + asked_for, + ); + // The index is on this generation's own columns, under the + // names they had when it was sealed. Every queried column has + // to be resolved: a rename moved the table's name while the + // file still holds the old one. + let mut stored_columns = Vec::with_capacity(columns.len()); + for column in columns { + let Some(stored) = generation.stored_name(column) else { + // Sealed before the column existed, so it has nothing to + // match -- and nothing to give a predicate spanning it. + return self.empty_plan(target_schema); + }; + stored_columns.push((column.clone(), stored.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 (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 // 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); } - scanner.full_text_search(bind(query)?)?; - scanner.create_plan().await + // A generation that stores every queried column under the name + // the table still uses needs no rebinding, so the tree reaches + // the scanner exactly as the other arms get it -- including a + // cross-column predicate, whose leaves carry their own bindings + // and must not be collapsed onto one column. + let bound_query = match stored_columns.as_slice() { + _ if stored_columns.iter().all(|(asked, stored)| asked == stored) => { + bind(query)? + } + // One column, moved: bind the whole tree to the name this + // generation stores it under. + [(_, stored)] => { + let bound = query.clone().with_column(stored.clone())?; + match limit.filter(|_| above.is_none()) { + Some(limit) => bound.limit(Some(limit as i64)), + None => bound.limit(None), + } + } + // Several columns, at least one of them renamed. Each leaf + // would need its own name and `with_column` rebinds the + // whole tree, collapsing a cross-column predicate onto one + // field -- which answers a different question. Refuse rather + // than rank on the wrong column silently. + moved => { + let renamed: Vec = moved + .iter() + .filter(|(asked, stored)| asked != stored) + .map(|(asked, stored)| format!("{asked} (stored as {stored})")) + .collect(); + return Err(Error::invalid_input(format!( + "a full-text search over several columns cannot read a sealed \ + generation in which one of them was renamed: {}. Compact the table \ + so the generation merges into base, or search the columns separately", + renamed.join(", ") + ))); + } + }; + 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), + None => Ok(reconciled), + } } LsmDataSource::ActiveMemTable { batch_store, @@ -1874,6 +1983,123 @@ mod tests { /// prefiltered candidate set must compose with cross-generation block-list /// filtering plus over-fetch. Gen 1's best predicate-matching hit (id=3) is /// superseded by gen 2; with over-fetch, gen 1 should still contribute id=4. + /// A cross-column predicate meets on rows, so every leaf has to reach its + /// own index. When a sealed generation stores one of those columns under an + /// older name, each leaf would need a different name -- and `with_column` + /// rebinds the whole tree, which would collapse the predicate onto one + /// field and answer a different question. Refused rather than ranked on the + /// wrong column. + #[tokio::test] + async fn a_cross_column_predicate_over_a_renamed_generation_is_refused() { + use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; + use crate::index::DatasetIndexExt; + use lance_core::datatypes::LANCE_FIELD_ID_KEY; + use lance_index::IndexType; + use lance_index::scalar::inverted::query::{BooleanQuery, MatchQuery, Occur}; + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let stored_schema = two_column_fts_schema(); + let tmp = tempfile::tempdir().unwrap(); + let base_uri = format!("{}/base", tmp.path().to_str().unwrap()); + let shard_id = uuid::Uuid::new_v4(); + + // The generation holds `title` and `body`, indexed under those names. + let gen1_uri = format!("{}/_mem_wal/{}/gen_1", base_uri, shard_id); + let mut gen1 = write_dataset( + &gen1_uri, + vec![make_two_column_batch( + &stored_schema, + &[(1, "alpha", "beta")], + )], + ) + .await; + for column in ["title", "body"] { + gen1.create_index( + &[column], + IndexType::Inverted, + Some(format!("{column}_fts")), + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + } + + // The table has since renamed `title` to `heading`. The id stays, which + // is what lets the generation be recognised as holding the same column. + let stamp = |name: &str, id: i32, nullable: bool| { + Field::new(name, DataType::Utf8, nullable).with_metadata(HashMap::from([( + LANCE_FIELD_ID_KEY.to_string(), + id.to_string(), + )])) + }; + let mut pk_meta = HashMap::from([(LANCE_FIELD_ID_KEY.to_string(), "0".to_string())]); + pk_meta.insert( + "lance-schema:unenforced-primary-key".to_string(), + "true".to_string(), + ); + let table_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(pk_meta), + stamp("heading", 1, true), + stamp("body", 2, true), + ])); + + let snapshot = ShardSnapshot::new(shard_id) + .with_current_generation(2) + .with_sstable(1, "gen_1".to_string()); + let collector = + LsmDataSourceCollector::without_base_table(base_uri.clone(), vec![snapshot]); + let planner = + LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], table_schema.clone()) + .with_identity_schema(Arc::clone(&table_schema)); + + // One predicate over both fields, under the names the table uses now. + let query = |first: &str| { + FullTextSearchQuery::new_query(IndexFtsQuery::Boolean(BooleanQuery::new(vec![ + ( + Occur::Must, + MatchQuery::new("alpha".to_string()) + .with_column(Some(first.to_string())) + .into(), + ), + ( + Occur::Must, + MatchQuery::new("beta".to_string()) + .with_column(Some("body".to_string())) + .into(), + ), + ]))) + }; + + let error = planner + .plan_search(query("heading"), Some(10), None) + .await + .expect_err("a renamed column in a cross-column predicate must be refused"); + let message = error.to_string(); + assert!( + message.contains("renamed") && message.contains("heading"), + "the refusal should name the column that moved, got: {message}" + ); + + // The control: the same generation, the same predicate, but the table + // still calls the column what the generation does. Nothing to rebind, + // so the tree reaches the scanner and the arm plans. + let unmoved = Arc::new(ArrowSchema::new(vec![ + table_schema.field(0).clone(), + stamp("title", 1, true), + stamp("body", 2, true), + ])); + let snapshot = ShardSnapshot::new(shard_id) + .with_current_generation(2) + .with_sstable(1, "gen_1".to_string()); + let collector = LsmDataSourceCollector::without_base_table(base_uri, vec![snapshot]); + LsmFtsSearchPlanner::new(collector, vec!["id".to_string()], unmoved.clone()) + .with_identity_schema(unmoved) + .plan_search(query("title"), Some(10), None) + .await + .expect("an unmoved cross-column predicate still plans over the generation"); + } + #[tokio::test] async fn prefilter_on_sstable_composes_with_block_list() { use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot; 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..4e93171f12b --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/scanner/generation.rs @@ -0,0 +1,619 @@ +// 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, without_field_id}; +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 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, + /// 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, also carrying field ids. + table_schema: SchemaRef, + pk_columns: 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 { + /// `projection` is what the caller asks for, under the table's names. + pub(super) fn new( + dataset_schema: &lance_core::datatypes::Schema, + table_schema: &SchemaRef, + pk_columns: &[String], + projection: Vec, + ) -> Self { + 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_schema, + names, + table_schema, + pk_columns, + projection, + } + } + + /// 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(|(_, 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.projection.iter().any(|w| w == column) { + self.projection.push(column.to_string()); + } + } + + /// 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.projection + .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_schema.field_with_name(name).ok()) + .flatten() + .map(|f| f.name().as_str()) + } + + /// `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 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.stored_as_declared(&c.name)); + if !pushable { + return None; + } + expr.clone() + .transform(|e| match e { + Expr::Column(mut c) => { + // `stored_name` is total over the refs, checked above. + c.name = self.stored_name(&c.name).expect("checked").to_string(); + Ok(Transformed::yes(Expr::Column(c))) + } + other => Ok(Transformed::no(other)), + }) + .map(|t| t.data) + .ok() + } + + /// 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 stored_as_declared(&self, column: &str) -> bool { + let Some(stored_column) = self.stored_name(column) else { + return false; + }; + let (Ok(stored_field), Ok(declared)) = ( + self.stored_schema.field_with_name(stored_column), + self.table_schema.field_with_name(column), + ) else { + return false; + }; + // 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 + /// 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.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() { + return Ok(scan); + } + Ok(Arc::new(ReconcileExec::new(scan, Arc::new(plan)))) + } + + /// `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 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 => without_field_id(field), + }) + .collect(); + Schema::new_with_metadata(fields, source.metadata().clone()) + } + + /// 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 + /// 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 + .projection + .iter() + .filter_map(|name| { + 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_column| source.field_with_name(stored_column).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. + 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 that could not be pushed into the +/// generation's own scan. +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!("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}")))?; + Ok(Arc::new( + FilterExec::try_new(physical, plan).map_err(|e| Error::internal(format!("filter: {e}")))?, + )) +} + +/// 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: &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_schema + .fields() + .iter() + .filter(|f| f.name() != TOMBSTONE && !is_system_column(f.name())) + .filter(|f| table_schema.field_with_name(f.name()).is_ok()) + .map(|f| (f.name().clone(), f.name().clone())) + .collect(); + } + stored_schema + .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) -> 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, + )) + } + (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, + } + } + let fields: Vec = schema + .fields() + .iter() + .map(|field| restore(field, stored_schema.fields())) + .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::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 + /// 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_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_schema, + &["id".to_string()], + projection.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"], + ) + } + + /// 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"]); + 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( + 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_eq!( + read.to_stored(&col("amount").eq(lit(1i64))), + Some(col("value").eq(lit(1i64))), + ); + } + + #[test] + fn a_predicate_naming_a_column_that_did_not_move_is_left_alone() { + let read = renamed(); + let expr = col("id").eq(lit(1i64)); + 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"], + ); + assert_eq!(read.to_stored(&col("added").eq(lit(1i64))), 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 = Schema::new(vec![ + with_id("id", DataType::Int64, 0), + with_id("value", DataType::Int64, 1), + with_id(TOMBSTONE, DataType::Boolean, 2), + ]); + let table_schema = Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Int64, true), + ]); + 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 2a83447fb17..061e72971b8 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/planner.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/planner.rs @@ -18,6 +18,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::generation::{GenerationRead, filter_above}; use super::projection::{ build_scanner_projection, canonical_output_schema, null_columns, project_to_canonical, validate_projection_names, @@ -45,6 +46,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. @@ -61,11 +66,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, @@ -185,9 +192,13 @@ impl LsmScanPlanner { (Some(n), false, false) => Some(n), _ => None, }; - let scan = 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 @@ -228,12 +239,39 @@ 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. `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. 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.first()) + .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 +281,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,34 +370,58 @@ impl LsmScanPlanner { .await?; let mut scanner = dataset.scan(); - let cols = + // 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 asked_for = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; + let mut generation = GenerationRead::new( + dataset.schema(), + &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 + // 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, above) = generation.split_filter(filter); + 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. + // 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(filter); + folded = fold_not_tombstone(stored_filter.as_ref()); Some(&folded) } else { - filter + stored_filter.as_ref() }; 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. - if let Some(fetch) = fetch { + // A limit under a filter that has not run would cut rows the + // filter never saw. + if let Some(fetch) = fetch.filter(|_| above.is_none()) { scanner.limit(Some(fetch as i64), None)?; } - 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), + } } LsmDataSource::ActiveMemTable { batch_store, @@ -372,6 +434,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 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. 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/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index 9aea4e9f4e7..2bd0c3a910d 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -34,11 +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::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable}; +use crate::dataset::mem_wal::reconcile::relabel_to; 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 { @@ -262,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))); @@ -683,13 +700,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(), + &self.identity_schema, + &self.pk_columns, + 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 +1042,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/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index fd52ac90c0b..6e83dae547c 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,33 @@ 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)) + } + // 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(), + 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 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!( + "no null literal for column '{name}' of type {}: {e}", + field.data_type() + )) + })?, + )), }; project_exprs.push((expr, name.clone())); } 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..61ae36bd854 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs @@ -28,6 +28,7 @@ use crate::io::exec::TakeExec; use super::collector::LsmDataSourceCollector; use super::data_source::LsmDataSource; +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, @@ -81,6 +82,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.). @@ -133,6 +137,7 @@ impl LsmVectorSearchPlanner { Self { collector, pk_columns, + identity_schema: base_schema.clone(), base_schema, vector_column, distance_type, @@ -174,6 +179,16 @@ impl LsmVectorSearchPlanner { self } + /// 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 + } + /// Set the session used to open SSTables. pub fn with_session(mut self, session: Arc) -> Self { self.session = Some(session); @@ -313,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?; @@ -495,18 +515,50 @@ 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 asked_for = build_scanner_projection(projection, &self.base_schema, &self.pk_columns); - scanner.project(&cols.iter().map(|s| s.as_str()).collect::>())?; - if let Some(ref filter) = self.filter { + let mut generation = GenerationRead::new( + dataset.schema(), + &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. + 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. + 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 // 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)?; + // 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(_) => 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); scanner.nprobes(nprobes); scanner.distance_metric(self.distance_type); @@ -514,7 +566,14 @@ impl LsmVectorSearchPlanner { scanner.ef(ef); } scanner.fast_search(); - scanner.create_plan().await + // 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), + None => 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 d0d1246ea83..8167f55aa06 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -18,8 +18,9 @@ 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_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; @@ -1389,6 +1390,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 +1436,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,27 +1602,62 @@ 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. +/// 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) +} + +/// Re-label `batch` to the storage schema, matching columns by **field id** +/// where both sides carry one, and by **name** otherwise. /// -/// 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 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, + 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 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 inject _tombstone column (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 @@ -2056,6 +2095,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` 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, @@ -2090,8 +2135,10 @@ 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); + 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(); let shard_id = config.shard_id; @@ -2371,6 +2418,7 @@ impl ShardWriter { manifest, manifest.current_generation, make_bound_memtable, + &pk_columns, &flusher, &wal_flusher, index_configs, @@ -2639,7 +2687,7 @@ impl ShardWriter { // `_tombstone`. let batches = batches .into_iter() - .map(|b| ensure_tombstone_column(b, &writer_state.schema)) + .map(|b| conform_live_batch(b, &writer_state.schema, &writer_state.pk_columns)) .collect::>>()?; self.put_memtable(batches, state, writer_state, backpressure) .await @@ -2764,7 +2812,7 @@ impl ShardWriter { // Mirrors `put`. let batches = batches .into_iter() - .map(|b| ensure_tombstone_column(b, &writer_state.schema)) + .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 @@ -4530,8 +4578,10 @@ mod tests { use super::*; use crate::dataset::mem_wal::test_util::failing_memory_store; use arrow_array::{FixedSizeListArray, Float32Array, Int32Array, 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; @@ -4651,10 +4701,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 +4718,251 @@ 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 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_refuses_a_column_whose_type_moved() { + 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::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec!["a"])), + Arc::new(Int32Array::from(vec![7])), + ], + ) + .unwrap(); + + 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. + #[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 struct column renamed at the top level is matched by the id on the + /// column itself, and taken whole. + /// + /// 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 { + 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] + 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] fn test_build_tombstone_batch_shape() { let storage = schema_with_tombstone(&create_test_schema()); diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 8c1977d1605..07bb80d780e 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,50 @@ pub(super) async fn add_columns( .await } +/// What `alter_columns` refuses on a table with a MemWAL, and why. +/// +/// 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, + /// 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, +} + +/// 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(()); + }; + if dataset.mem_wal_index_details().await?.is_none() { + return Ok(()); + } + 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]) { let Some(first_fragment) = fragments.first() else { return; @@ -729,6 +774,15 @@ pub(super) async fn alter_columns( dataset: &mut Dataset, alterations: &[ColumnAlteration], ) -> Result<()> { + 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. let mut new_schema = dataset.schema().clone();