fix(mem_wal): reconcile a schema change through one plan

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,
 };

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<ArrowField> = 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<ArrowField> = 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<ArrowSchema> {
     ))
 }
 
+/// `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<ArrowSchema>,
+    pk_columns: &[String],
+    batches: Vec<RecordBatch>,
+) -> lance_core::Result<Vec<RecordBatch>> {
+    // 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::<Vec<_>>(),
+        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::<arrow_array::BooleanArray>()
+            .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::<arrow_array::StructArray>()
+            .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::<arrow_array::Int64Array>()
+            .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::<arrow_array::Int64Array>()
+            .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()]);