feat(mem_wal): support delete against non-nullable base columns

rust/lance/src/dataset/mem_wal.rs

@@ -58,13 +58,13 @@ use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
 /// its primary key, carrying null in every non-PK column, that wins
 /// newest-per-PK resolution and is then silently dropped from query results.
 ///
-/// The column is owned end-to-end by lance: callers pass the base schema and
+/// The column is owned end-to-end by lance: callers pass the logical schema and
 /// lance injects the column on the write path ([`write::ShardWriter::put`] /
 /// [`write::ShardWriter::delete`]), so no caller ever constructs or names it.
 pub const TOMBSTONE: &str = "_tombstone";
 
-/// The mem_wal tombstone field appended to the base schema to form the
-/// memtable/generation schema.
+/// The mem_wal tombstone field appended to the logical schema on the way to the
+/// storage schema.
 ///
 /// Non-nullable: the write path always populates it (`false` for normal rows,
 /// `true` for tombstones). Non-nullability also lets the point-lookup base arm
@@ -74,8 +74,48 @@ pub fn tombstone_field() -> ArrowField {
     ArrowField::new(TOMBSTONE, DataType::Boolean, false)
 }
 
-/// Extend a base schema with the trailing `_tombstone` column to form the
-/// mem_wal memtable/generation schema.
+/// Derive a shard's *storage* schema from its *logical* (base table) schema by
+/// widening every top-level field to nullable except the primary key and
+/// `_tombstone`.
+///
+/// A tombstone carries the primary key and null in every other column, so the
+/// memtable, WAL entries, and SSTables must permit a null wherever the base
+/// table does not. The logical schema stays the caller's contract:
+/// [`write::ShardWriter::put`] validates against it and the scan path narrows
+/// back to it.
+///
+/// Top-level only — Arrow validates nullability only at the top level of a
+/// `RecordBatch`, so a vector column's item field gains no validity layer.
+/// Primary keys are excluded because [`lance_core::datatypes::Schema`] requires
+/// them to be non-nullable and a tombstone always carries a real key;
+/// `_tombstone` because the write path always populates it. Idempotent.
+pub fn relax_non_pk_nullability(
+    logical_schema: &ArrowSchema,
+    pk_columns: &[String],
+) -> Arc<ArrowSchema> {
+    let fields: Vec<ArrowField> = logical_schema
+        .fields()
+        .iter()
+        .map(|field| {
+            let keep = field.is_nullable()
+                || field.name() == TOMBSTONE
+                || pk_columns.iter().any(|c| c == field.name());
+            let field = field.as_ref().clone();
+            if keep {
+                field
+            } else {
+                field.with_nullable(true)
+            }
+        })
+        .collect();
+    Arc::new(ArrowSchema::new_with_metadata(
+        fields,
+        logical_schema.metadata().clone(),
+    ))
+}
+
+/// Extend the logical schema with the trailing `_tombstone` column — the
+/// intermediate [`relax_non_pk_nullability`] widens into the storage schema.
 ///
 /// Idempotent: a schema that already carries `_tombstone` (a reopen/replay
 /// path) is returned unchanged. Schema-level metadata and per-field metadata
@@ -104,3 +144,94 @@ pub use wal::{BatchDurableWatcher, WalAppendResult, WalAppender, WalReadEntry, W
 pub use write::ShardWriter;
 pub use write::ShardWriterConfig;
 pub use write::WriteResult;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow_schema::Fields;
+
+    fn logical() -> ArrowSchema {
+        ArrowSchema::new(vec![
+            ArrowField::new("id", DataType::Int32, false),
+            ArrowField::new("count", DataType::Int64, false),
+            ArrowField::new("note", DataType::Utf8, true),
+        ])
+    }
+
+    #[test]
+    fn relax_widens_every_non_pk_field_and_leaves_the_key_alone() {
+        let relaxed = relax_non_pk_nullability(&logical(), &["id".to_string()]);
+
+        assert!(
+            !relaxed.field(0).is_nullable(),
+            "the primary key stays strict"
+        );
+        assert!(
+            relaxed.field(1).is_nullable(),
+            "`count` must accept a tombstone null"
+        );
+        assert!(
+            relaxed.field(2).is_nullable(),
+            "already-nullable is untouched"
+        );
+    }
+
+    #[test]
+    fn relax_leaves_nested_fields_exactly_as_declared() {
+        // Arrow validates nullability only at the top level, and a vector
+        // column's item field must not gain a validity layer.
+        let item = Arc::new(ArrowField::new("item", DataType::Float32, false));
+        let child = ArrowField::new("a", DataType::Int32, false);
+        let schema = ArrowSchema::new(vec![
+            ArrowField::new("id", DataType::Int32, false),
+            ArrowField::new("vector", DataType::FixedSizeList(item, 4), false),
+            ArrowField::new("s", DataType::Struct(Fields::from(vec![child])), false),
+        ]);
+
+        let relaxed = relax_non_pk_nullability(&schema, &["id".to_string()]);
+
+        assert!(relaxed.field(1).is_nullable());
+        match relaxed.field(1).data_type() {
+            DataType::FixedSizeList(f, _) => assert!(!f.is_nullable(), "item field untouched"),
+            other => panic!("expected FixedSizeList, got {other:?}"),
+        }
+        match relaxed.field(2).data_type() {
+            DataType::Struct(fields) => assert!(!fields[0].is_nullable(), "child field untouched"),
+            other => panic!("expected Struct, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn relax_keeps_tombstone_non_nullable_and_is_idempotent() {
+        let pk = ["id".to_string()];
+        let once = relax_non_pk_nullability(&schema_with_tombstone(&logical()), &pk);
+        let twice = relax_non_pk_nullability(&once, &pk);
+
+        let tombstone = once.field_with_name(TOMBSTONE).unwrap();
+        assert!(
+            !tombstone.is_nullable(),
+            "the write path always populates _tombstone"
+        );
+        assert_eq!(once, twice);
+    }
+
+    #[test]
+    fn relax_preserves_schema_and_field_metadata() {
+        // The `lance-schema:unenforced-primary-key` marker rides on field
+        // metadata, so losing it here would silently drop the shard's PK.
+        let marked = ArrowField::new("count", DataType::Int64, false)
+            .with_metadata([("k".to_string(), "v".to_string())].into());
+        let schema = ArrowSchema::new_with_metadata(
+            vec![ArrowField::new("id", DataType::Int32, false), marked],
+            [("s".to_string(), "m".to_string())].into(),
+        );
+
+        let relaxed = relax_non_pk_nullability(&schema, &["id".to_string()]);
+
+        assert_eq!(relaxed.metadata().get("s").map(String::as_str), Some("m"));
+        assert_eq!(
+            relaxed.field(1).metadata().get("k").map(String::as_str),
+            Some("v")
+        );
+    }
+}

rust/lance/src/dataset/mem_wal/scanner/exec.rs

@@ -10,12 +10,14 @@
 //! - [`BloomFilterGuardExec`]: Guards child execution with bloom filter check
 //! - [`CoalesceFirstExec`]: Returns first non-empty result with short-circuit
 //! - [`PkBlockFilterExec`]: Drops rows whose PK was superseded by a newer generation (the cross-generation block-list)
+//! - [`SchemaRelabelExec`]: Re-labels batches to an exact schema (the logical/storage nullability boundary)
 
 mod bloom_guard;
 mod coalesce_first;
 mod generation_tag;
 mod pk;
 mod pk_block_filter;
+mod schema_relabel;
 
 pub use bloom_guard::{BloomFilterGuardExec, compute_pk_hash_from_scalars};
 pub use coalesce_first::CoalesceFirstExec;
@@ -25,3 +27,4 @@ pub use pk::{
     validate_pk_types,
 };
 pub use pk_block_filter::PkBlockFilterExec;
+pub use schema_relabel::SchemaRelabelExec;