diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 784ee1fa76c..40e7dbde46d 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/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 { + let fields: Vec = 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") + ); + } +} diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/scanner/exec.rs index 1498c5f60ea..9c47c893d8d 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec.rs +++ b/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; 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 new file mode 100644 index 00000000000..33131c95d5e --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Schema re-labeling execution node. + +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, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt}; + +/// Re-labels every batch to an exact target schema, leaving the arrays +/// untouched. +/// +/// A shard's storage schema widens non-PK columns to nullable (see +/// `relax_non_pk_nullability`), so plan arms disagree with the base-table arm on +/// nullability alone. `ProjectionExec` cannot pin this down: DataFusion derives +/// its output nullability from the expressions, not from the schema the planner +/// intended. +/// +/// Used both ways: **widening**, so arms agree before `UnionExec` / +/// `CoalesceFirstExec`; and **narrowing** at the scan's output boundary, back to +/// the logical schema. Narrowing doubles as the tombstone-leak check — +/// `RecordBatch::try_new` rejects a null in a non-nullable column, so a leak +/// errors instead of reaching the caller. +#[derive(Debug)] +pub struct SchemaRelabelExec { + input: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl SchemaRelabelExec { + /// Wrap `input` so its batches are re-labeled to `schema`, which the caller + /// must keep column-compatible (same count, order, and data types); only + /// names, nullability, and metadata may differ. A mismatch surfaces per + /// batch at execution time, not at plan time. + pub fn new(input: Arc, schema: SchemaRef) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Self { + input, + schema, + properties, + } + } +} + +impl DisplayAs for SchemaRelabelExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => { + write!(f, "SchemaRelabelExec") + } + } + } +} + +impl ExecutionPlan for SchemaRelabelExec { + fn name(&self) -> &str { + "SchemaRelabelExec" + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DFResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "SchemaRelabelExec requires exactly one child".to_string(), + )); + } + Ok(Arc::new(Self::new( + children[0].clone(), + self.schema.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DFResult { + Ok(Box::pin(SchemaRelabelStream { + input: self.input.execute(partition, context)?, + schema: self.schema.clone(), + })) + } +} + +struct SchemaRelabelStream { + input: SendableRecordBatchStream, + schema: SchemaRef, +} + +impl Stream for SchemaRelabelStream { + 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))) => { + let schema = self.schema.clone(); + // Guards the column-less case: `try_new` infers the row count + // from the first column and errors when there is none. + let relabeled = if batch.num_rows() == 0 { + Ok(RecordBatch::new_empty(schema)) + } else { + RecordBatch::try_new(schema, batch.columns().to_vec()) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) + }; + Poll::Ready(Some(relabeled)) + } + other => other, + } + } +} + +impl datafusion::physical_plan::RecordBatchStream for SchemaRelabelStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Int32Array, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::prelude::SessionContext; + use datafusion_physical_plan::test::TestMemoryExec; + use futures::TryStreamExt; + + fn schema_with(nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, nullable), + ])) + } + + fn source(batch: RecordBatch) -> Arc { + TestMemoryExec::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap() + } + + fn batch(schema: SchemaRef, names: Vec>) -> RecordBatch { + let ids: Vec = (0..names.len() as i32).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + ], + ) + .unwrap() + } + + async fn run(plan: Arc) -> DFResult> { + let ctx = SessionContext::new(); + plan.execute(0, ctx.task_ctx())?.try_collect().await + } + + #[tokio::test] + async fn widening_preserves_rows_and_reports_target_schema() { + let input = source(batch(schema_with(false), vec![Some("a"), Some("b")])); + let relabeled = Arc::new(SchemaRelabelExec::new(input, schema_with(true))); + + assert_eq!(relabeled.schema(), schema_with(true)); + let out = run(relabeled).await.unwrap(); + assert_eq!(out.len(), 1); + assert_eq!(out[0].schema(), schema_with(true)); + assert_eq!(out[0].num_rows(), 2); + } + + #[tokio::test] + async fn narrowing_succeeds_when_no_nulls_remain() { + // Post-tombstone-filter: `name` is nullable in storage, but every + // surviving row has a value. + let input = source(batch(schema_with(true), vec![Some("a"), Some("b")])); + let relabeled = Arc::new(SchemaRelabelExec::new(input, schema_with(false))); + + let out = run(relabeled).await.unwrap(); + assert_eq!(out[0].schema(), schema_with(false)); + assert_eq!(out[0].num_rows(), 2); + } + + #[tokio::test] + async fn narrowing_rejects_a_surviving_null() { + // A tombstone that escaped its filter must error here, not reach the + // caller as a row of nulls. + let input = source(batch(schema_with(true), vec![Some("a"), None])); + let relabeled = Arc::new(SchemaRelabelExec::new(input, schema_with(false))); + + let error = run(relabeled).await.unwrap_err().to_string(); + assert!( + error.contains("non-nullable") && error.contains("name"), + "expected a nullability error naming the column, got: {error}" + ); + } + + #[tokio::test] + async fn empty_batch_is_relabeled_without_row_count_inference() { + let input = source(batch(schema_with(true), vec![])); + let relabeled = Arc::new(SchemaRelabelExec::new(input, schema_with(false))); + + let out = run(relabeled).await.unwrap(); + assert!(out.iter().all(|b| b.num_rows() == 0)); + assert!(out.iter().all(|b| b.schema() == schema_with(false))); + } +} 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 9fc44290209..fd9daa73756 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -27,15 +27,15 @@ use lance_core::{Result, is_system_column}; use lance_datafusion::exec::OneShotExec; use tracing::instrument; -use crate::dataset::mem_wal::TOMBSTONE; use crate::dataset::mem_wal::index::IndexStore; use crate::dataset::mem_wal::memtable::batch_store::BatchStore; +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::projection::{ - DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns, + 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}; @@ -321,7 +321,7 @@ impl LsmPointLookupPlanner { ) -> Result { let canonical = canonical_output_schema(projection, &self.base_schema, &self.pk_columns, false); - let target = carry_schema(&canonical); + let target = carry_schema(&canonical, &self.pk_columns); let mut out: Vec = Vec::with_capacity(keys.len()); for key in keys { if let Some(b) = self.lookup_keep_tombstone(key, projection).await? { @@ -728,7 +728,7 @@ impl LsmPointLookupPlanner { // Output carries `_tombstone` (canonical + the marker) so it survives // the union/coalesce to the post-coalesce filter; base / legacy sources // that lack the column get a synthesized `false`. - project_to_carry(scan, &target) + project_to_carry(scan, &target, &self.pk_columns) } /// Create an empty execution plan with the canonical output schema. @@ -756,11 +756,18 @@ fn cols_with_tombstone(cols: &[String], present: bool) -> Vec { out } -/// Carry schema = canonical output + a trailing non-nullable `_tombstone` -/// Boolean. Non-nullable so the base arm's synthesized `Literal(false)` matches -/// the WAL arms' real column under `CoalesceFirstExec`'s exact-schema check. -fn carry_schema(canonical: &SchemaRef) -> SchemaRef { - let mut fields: Vec> = canonical.fields().iter().cloned().collect(); +/// Carry schema = canonical output at the storage schema's nullability, plus a +/// trailing non-nullable `_tombstone` Boolean. +/// +/// Widened because tombstone rows — null in every non-PK column — are still in +/// flight here; [`filter_tombstones_after_coalesce`] drops them on the far side +/// of `CoalesceFirstExec`, and the narrowing back to the logical schema happens +/// after that. `_tombstone` stays non-nullable so the base arm's synthesized +/// `Literal(false)` matches the WAL arms' real column under +/// `CoalesceFirstExec`'s exact-schema check. +fn carry_schema(canonical: &SchemaRef, pk_columns: &[String]) -> SchemaRef { + let widened = relax_non_pk_nullability(canonical, pk_columns); + let mut fields: Vec> = widened.fields().iter().cloned().collect(); fields.push(Arc::new(Field::new(TOMBSTONE, DataType::Boolean, false))); Arc::new(Schema::new(fields)) } @@ -772,9 +779,10 @@ fn carry_schema(canonical: &SchemaRef) -> SchemaRef { fn project_to_carry( plan: Arc, canonical: &SchemaRef, + pk_columns: &[String], ) -> Result> { let input = plan.schema(); - let carry = carry_schema(canonical); + let carry = carry_schema(canonical, pk_columns); let mut project_exprs: Vec<(Arc, String)> = Vec::with_capacity(carry.fields().len()); for field in carry.fields() { @@ -798,11 +806,11 @@ fn project_to_carry( }; project_exprs.push((expr, name.clone())); } - Ok(Arc::new( - ProjectionExec::try_new(project_exprs, plan).map_err(|e| { - lance_core::Error::internal(format!("Failed to build carry ProjectionExec: {}", e)) - })?, - )) + let projected = Arc::new(ProjectionExec::try_new(project_exprs, plan).map_err(|e| { + lance_core::Error::internal(format!("Failed to build carry ProjectionExec: {}", e)) + })?); + // `CoalesceFirstExec` panics unless every arm lands on exactly `carry`. + Ok(force_schema(projected, &carry)) } /// Drop tombstone rows after `CoalesceFirstExec` has already picked the newest diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index a39c83a8f1b..48c0d655496 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs @@ -25,6 +25,8 @@ 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; + /// Column name for distance in vector search results. pub const DISTANCE_COLUMN: &str = "_distance"; @@ -189,9 +191,28 @@ pub fn null_columns( Ok(Arc::new(projection_exec)) } +/// Force `plan` to report exactly `target_schema`; a no-op when they agree. +/// +/// A `ProjectionExec` cannot do this on its own — DataFusion derives its field +/// nullability from the expressions, not from the schema the planner asked for +/// — and the storage schema's widened non-PK columns leave the WAL arms +/// disagreeing with the base arm, which `CoalesceFirstExec` and +/// `concat_batches` both reject. +pub(super) fn force_schema( + plan: Arc, + target_schema: &SchemaRef, +) -> Arc { + if plan.schema() == *target_schema { + return plan; + } + Arc::new(SchemaRelabelExec::new(plan, target_schema.clone())) +} + /// Wrap `plan` to emit exactly `target_schema`. Source columns are /// forwarded by name; system / `_distance` cols missing from the source /// are NULL-filled. Other missing columns are an internal error. +/// +/// Reports `target_schema` exactly, nullability included — see [`force_schema`]. pub fn project_to_canonical( plan: Arc, target_schema: &SchemaRef, @@ -222,7 +243,7 @@ pub fn project_to_canonical( let projection_exec = ProjectionExec::try_new(project_exprs, plan).map_err(|e| { lance_core::Error::internal(format!("Failed to build canonical ProjectionExec: {}", e)) })?; - Ok(Arc::new(projection_exec)) + Ok(force_schema(Arc::new(projection_exec), target_schema)) } #[cfg(test)] diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 7f05cbfff3e..2b360b3633a 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -53,7 +53,7 @@ use super::wal::{ BatchDurableWatcher, TriggerIndexApply, TriggerWalFlush, WalAppender, WalFlushSource, WalOnlyState, WalRetryConfig, WalTailer, WriterCursors, apply_index_range, empty_flush_result, }; -use super::{TOMBSTONE, schema_with_tombstone}; +use super::{TOMBSTONE, relax_non_pk_nullability, schema_with_tombstone}; use crate::session::Session; use super::manifest::ShardManifestStore; @@ -902,14 +902,14 @@ async fn replay_memtable_from_wal( // Fence sentinels deserialize to zero batches and are skipped // here — they carry only a position, no rows. if !entry.batches.is_empty() { - // Entries written before deletes existed lack `_tombstone`; - // inject `false` so they match the extended memtable schema. - // Normal entries already carry it and pass through unchanged. - let target_schema = active.schema().clone(); + // Re-label every replayed entry to the current storage + // schema; entries written before deletes existed also need + // `_tombstone = false` injected. + let storage_schema = active.schema().clone(); let batches = entry .batches .into_iter() - .map(|b| ensure_tombstone_column(b, &target_schema)) + .map(|b| ensure_tombstone_column(b, &storage_schema)) .collect::>>()?; // Seal + flush at the entry boundary on the *same* criteria the @@ -1045,46 +1045,45 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// Ensure `batch` carries the `_tombstone` column required by the extended -/// memtable schema, injecting `false` for every row when it is absent. +/// Re-label `batch` to the storage schema, injecting `_tombstone = false` when +/// the column is absent — callers pass logical-shaped batches, and WAL entries +/// written before deletes existed lack it. /// -/// Used on the normal write path ([`ShardWriter::put`]) where callers pass -/// base-shaped batches, and on WAL replay of entries written before deletes -/// existed (legacy entries lack the column). A batch that already carries -/// `_tombstone` (a normal replayed entry) is returned unchanged. +/// A batch that already carries `_tombstone` is re-labeled rather than passed +/// through, so an entry written under an older storage schema replays into the +/// current one. fn ensure_tombstone_column( batch: RecordBatch, - target_schema: &Arc, + storage_schema: &Arc, ) -> Result { - if batch.schema().column_with_name(TOMBSTONE).is_some() { - return Ok(batch); - } let n = batch.num_rows(); let mut columns: Vec = batch.columns().to_vec(); - columns.push(Arc::new(BooleanArray::from(vec![false; n]))); - RecordBatch::try_new(target_schema.clone(), columns).map_err(|e| { + if batch.schema().column_with_name(TOMBSTONE).is_none() { + columns.push(Arc::new(BooleanArray::from(vec![false; n]))); + } + RecordBatch::try_new(storage_schema.clone(), columns).map_err(|e| { Error::invalid_input(format!( - "failed to inject _tombstone column (does the batch match the base schema?): {}", + "failed to inject _tombstone column (does the batch match the base table schema?): {}", e )) }) } -/// Build a tombstone batch from a key-only `keys` batch: the primary key -/// columns are carried through, `_tombstone` is set to `true`, and every other -/// column in the memtable schema is null. +/// Build a tombstone batch from a key-only `keys` batch: primary keys carried +/// through, `_tombstone` true, every other column null. /// -/// Errors if `keys` is missing a primary key column, or if a non-PK column is -/// non-nullable (a tombstone must null it) — surfaced via the `RecordBatch` -/// validation. +/// The storage schema's non-PK columns are nullable whatever the base table +/// declares — that is what lets a table with non-nullable columns have +/// tombstones at all. Primary keys stay non-nullable, so the `RecordBatch` +/// validation below still rejects a null, mistyped, or missing key. fn build_tombstone_batch( keys: &RecordBatch, - target_schema: &Arc, + storage_schema: &Arc, pk_columns: &[String], ) -> Result { let n = keys.num_rows(); - let mut columns: Vec = Vec::with_capacity(target_schema.fields().len()); - for field in target_schema.fields() { + let mut columns: Vec = Vec::with_capacity(storage_schema.fields().len()); + for field in storage_schema.fields() { let name = field.name(); if name == TOMBSTONE { columns.push(Arc::new(BooleanArray::from(vec![true; n]))); @@ -1100,9 +1099,9 @@ fn build_tombstone_batch( columns.push(new_null_array(field.data_type(), n)); } } - RecordBatch::try_new(target_schema.clone(), columns).map_err(|e| { + RecordBatch::try_new(storage_schema.clone(), columns).map_err(|e| { Error::invalid_input(format!( - "failed to build tombstone batch (is every non-primary-key column nullable?): {}", + "failed to build tombstone batch (do the delete keys match the primary key?): {}", e )) }) @@ -1501,6 +1500,12 @@ pub struct ShardWriter { manifest_store: Arc, stats: SharedWriteStats, mode: WriterMode, + /// The base table's schema as the caller passed it — no `_tombstone`, + /// nullability untouched. Caller input is held to this and the scan path + /// narrows back to it; the memtable, WAL, and SSTables carry the widened + /// storage schema ([`relax_non_pk_nullability`]) instead. See + /// [`Self::validate_against_logical_schema`]. + logical_schema: Arc, } impl ShardWriter { @@ -1539,10 +1544,12 @@ impl ShardWriter { )); } - // Callers pass the base schema; lance owns the `_tombstone` column and - // appends it here so the memtable/generation schema = base + tombstone. - // Idempotent, so a reopen that already extended the schema is a no-op. - let schema = schema_with_tombstone(&schema); + // The caller's schema is the shard's logical schema; the storage schema + // is derived from it below, once the primary key is known. lance owns + // `_tombstone` and appends it here — idempotent, so a reopen that + // already extended the schema is a no-op. + let logical_schema = schema; + let tombstoned = schema_with_tombstone(&logical_schema); let base_uri = base_uri.into(); let shard_id = config.shard_id; @@ -1560,7 +1567,7 @@ impl ShardWriter { // the schema) must fail here, before it can knock the healthy incumbent off // the shard. Memtable-only: WAL-only mode has no indexes to validate. let memtable_validation = if config.enable_memtable { - let lance_schema = Schema::try_from(schema.as_ref())?; + let lance_schema = Schema::try_from(tombstoned.as_ref())?; let pk_fields = lance_schema.unenforced_primary_key(); let pk_field_ids: Vec = pk_fields.iter().map(|f| f.id).collect(); let pk_columns: Vec = pk_fields.iter().map(|f| f.name.clone()).collect(); @@ -1569,9 +1576,18 @@ impl ShardWriter { // single row is accepted. Such a config fails deterministically on // every insert, including inserts replayed from the WAL — so once a row // is durable the shard can never reopen. Fail the open instead. - validate_index_configs(&index_configs, schema.as_ref(), &lance_schema, &pk_columns)?; + validate_index_configs( + &index_configs, + tombstoned.as_ref(), + &lance_schema, + &pk_columns, + )?; - Some((pk_field_ids, pk_columns)) + // Widen only now that the primary key is known — a tombstone nulls + // every non-PK column. `unenforced_primary_key` above ran against + // the logical schema, which is what enforces non-nullable PKs. + let storage_schema = relax_non_pk_nullability(&tombstoned, &pk_columns); + Some((pk_field_ids, pk_columns, storage_schema)) } else { None }; @@ -1632,11 +1648,11 @@ impl ShardWriter { let task_executor = Arc::new(TaskExecutor::new()); let mode = if config.enable_memtable { - let (pk_field_ids, pk_columns) = memtable_validation + let (pk_field_ids, pk_columns, storage_schema) = memtable_validation .expect("memtable_validation is Some when enable_memtable is true"); Self::open_memtable_mode( &config, - &schema, + &storage_schema, &manifest, &index_configs, pk_field_ids, @@ -1673,6 +1689,7 @@ impl ShardWriter { manifest_store, stats, mode, + logical_schema, }) } @@ -1954,6 +1971,7 @@ impl ShardWriter { #[instrument(name = "sw_put", level = "info", skip_all, fields(batch_count = batches.len(), shard_id = %self.config.shard_id))] pub async fn put(&self, batches: Vec) -> Result { Self::validate_non_empty(&batches)?; + self.validate_against_logical_schema(&batches)?; match &self.mode { WriterMode::MemTable { @@ -1961,9 +1979,8 @@ impl ShardWriter { writer_state, backpressure, } => { - // Inject `_tombstone = false` so the batch matches the - // extended memtable schema; callers only ever pass base-shaped - // batches and never name the column. + // Callers pass logical-shaped batches and never name + // `_tombstone`. let batches = batches .into_iter() .map(|b| ensure_tombstone_column(b, &writer_state.schema)) @@ -1992,9 +2009,8 @@ impl ShardWriter { /// its key: it wins newest-per-PK resolution (suppressing the older real /// row) and is then dropped from query results. /// - /// Only supported in memtable mode. Because a tombstone nulls every non-PK - /// column, those columns must be nullable in the base schema; a delete - /// against a schema with a non-nullable non-PK column errors. + /// Only supported in memtable mode. Works against non-nullable base columns: + /// tombstones live in the storage schema, which widens them to nullable. /// /// ``` /// # use lance::Result; @@ -2081,6 +2097,7 @@ impl ShardWriter { batches: Vec, ) -> Result<(WriteResult, Option)> { Self::validate_non_empty(&batches)?; + self.validate_against_logical_schema(&batches)?; match &self.mode { WriterMode::MemTable { @@ -2088,8 +2105,7 @@ impl ShardWriter { writer_state, backpressure, } => { - // Inject `_tombstone = false` to match the extended memtable - // schema, mirroring `put`. + // Mirrors `put`. let batches = batches .into_iter() .map(|b| ensure_tombstone_column(b, &writer_state.schema)) @@ -2103,6 +2119,31 @@ impl ShardWriter { } } + /// Reject caller input that violates the logical schema: wrong column count + /// or types, or a null in a column the base table declares non-nullable. + /// + /// The *only* gate on that contract. The storage schema no longer rejects a + /// caller's null, and nothing downstream does either — append and + /// `merge_insert` compare schemas with `NullabilityComparison::Ignore`, and + /// the encoder takes validity from the array, not the field — so a null that + /// gets past here reaches the base table silently. + /// + /// Runs before the WAL append: a batch rejected only after being appended + /// would fail identically on every replay, leaving the shard unable to + /// reopen. + fn validate_against_logical_schema(&self, batches: &[RecordBatch]) -> Result<()> { + for (i, batch) in batches.iter().enumerate() { + RecordBatch::try_new(self.logical_schema.clone(), batch.columns().to_vec()).map_err( + |e| { + Error::invalid_input(format!( + "batch {i} does not match the base table schema: {e}" + )) + }, + )?; + } + Ok(()) + } + fn validate_non_empty(batches: &[RecordBatch]) -> Result<()> { if batches.is_empty() { return Err(Error::invalid_input("Cannot write empty batch list")); @@ -3679,6 +3720,17 @@ mod tests { ])) } + /// [`create_pk_test_schema`] with a non-nullable `name` — the shape that + /// used to make `delete` fail. + fn create_strict_pk_test_schema() -> Arc { + let fields: Vec = create_pk_test_schema() + .fields() + .iter() + .map(|f| f.as_ref().clone().with_nullable(false)) + .collect(); + Arc::new(ArrowSchema::new(fields)) + } + fn id_only_keys(ids: &[i32]) -> RecordBatch { RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![Field::new( @@ -3694,9 +3746,9 @@ mod tests { #[test] fn test_ensure_tombstone_column_injects_false() { let base = create_test_schema(); - let target = schema_with_tombstone(&base); - let out = ensure_tombstone_column(create_test_batch(&base, 0, 3), &target).unwrap(); - assert_eq!(out.schema(), target); + let storage = schema_with_tombstone(&base); + let out = ensure_tombstone_column(create_test_batch(&base, 0, 3), &storage).unwrap(); + assert_eq!(out.schema(), storage); let ts = out .column_by_name(TOMBSTONE) .unwrap() @@ -3708,16 +3760,16 @@ mod tests { "put injects _tombstone = false" ); // Idempotent: a batch already carrying the column passes through. - let again = ensure_tombstone_column(out.clone(), &target).unwrap(); + let again = ensure_tombstone_column(out.clone(), &storage).unwrap(); assert_eq!(again.schema(), out.schema()); } #[test] fn test_build_tombstone_batch_shape() { - let target = schema_with_tombstone(&create_test_schema()); + let storage = schema_with_tombstone(&create_test_schema()); let tomb = - build_tombstone_batch(&id_only_keys(&[5, 7]), &target, &["id".to_string()]).unwrap(); - assert_eq!(tomb.schema(), target); + build_tombstone_batch(&id_only_keys(&[5, 7]), &storage, &["id".to_string()]).unwrap(); + assert_eq!(tomb.schema(), storage); assert_eq!(tomb.num_rows(), 2); let ids = tomb .column_by_name("id") @@ -3742,7 +3794,7 @@ mod tests { #[test] fn test_build_tombstone_batch_missing_pk_errors() { - let target = schema_with_tombstone(&create_test_schema()); + let storage = schema_with_tombstone(&create_test_schema()); let keys = RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![Field::new( "other", @@ -3752,18 +3804,62 @@ mod tests { vec![Arc::new(Int32Array::from(vec![1]))], ) .unwrap(); - assert!(build_tombstone_batch(&keys, &target, &["id".to_string()]).is_err()); + assert!(build_tombstone_batch(&keys, &storage, &["id".to_string()]).is_err()); } #[test] - fn test_build_tombstone_batch_non_nullable_nonpk_errors() { - // A tombstone must null every non-PK column; a non-nullable one fails. + fn test_build_tombstone_batch_nulls_non_nullable_base_column() { + // The point of the storage schema: a tombstone nulls `v` even though the + // base table declares it non-nullable. + let pk = ["id".to_string()]; let base = Arc::new(ArrowSchema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("v", DataType::Int32, false), ])); - let target = schema_with_tombstone(&base); - assert!(build_tombstone_batch(&id_only_keys(&[1]), &target, &["id".to_string()]).is_err()); + let storage = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); + + let batch = build_tombstone_batch(&id_only_keys(&[1]), &storage, &pk).unwrap(); + + assert!(batch["v"].is_null(0), "the tombstone must null `v`"); + assert!(!batch["id"].is_null(0), "the primary key survives"); + assert!( + batch[TOMBSTONE] + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + } + + #[test] + fn test_build_tombstone_batch_rejects_null_primary_key() { + // Primary keys are never relaxed, so the storage schema still rejects a + // null key — the delete path needs no separate check for it. + let pk = ["id".to_string()]; + let base = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let storage = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); + let keys = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])), + vec![Arc::new(Int32Array::from(vec![None::]))], + ) + .unwrap(); + + let error = build_tombstone_batch(&keys, &storage, &pk).unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected InvalidInput, got {error:?}" + ); + assert!( + error.to_string().contains("non-nullable"), + "error should name the nullability violation: {error}" + ); } #[tokio::test] @@ -3832,6 +3928,159 @@ mod tests { writer.close().await.unwrap(); } + /// Delete works against a base table with non-nullable non-PK columns, and + /// survivors come back through the narrowing egress relabel intact. + #[tokio::test] + async fn test_delete_against_non_nullable_base_column_round_trip() { + use crate::dataset::mem_wal::scanner::LsmScanner; + use futures::TryStreamExt; + + let (store, base_path, base_uri, _temp) = create_local_store().await; + let schema = create_strict_pk_test_schema(); + assert!( + !schema.field_with_name("name").unwrap().is_nullable(), + "the point of this test is a non-nullable non-PK column" + ); + let config = ShardWriterConfig { + shard_id: Uuid::new_v4(), + durable_write: true, + ..Default::default() + }; + let shard_id = config.shard_id; + let writer = ShardWriter::open( + store, + base_path, + base_uri.clone(), + config, + schema.clone(), + vec![], + ) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 5)]) + .await + .unwrap(); + writer.delete(vec![id_only_keys(&[2])]).await.unwrap(); + + let refs = writer.in_memory_memtable_refs().await.unwrap(); + let scanner = LsmScanner::without_base_table( + schema.clone(), + base_uri, + vec![], + vec!["id".to_string()], + ) + .with_in_memory_memtables(shard_id, refs); + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let mut rows: Vec<(i32, String)> = Vec::new(); + for b in &batches { + assert!( + !b.schema().field_with_name("name").unwrap().is_nullable(), + "egress must narrow back to the logical schema" + ); + let ids = b["id"].as_any().downcast_ref::().unwrap(); + let names = b["name"].as_any().downcast_ref::().unwrap(); + rows.extend((0..ids.len()).map(|i| (ids.value(i), names.value(i).to_string()))); + } + rows.sort_unstable(); + + assert_eq!( + rows, + vec![ + (0, "name_0".to_string()), + (1, "name_1".to_string()), + (3, "name_3".to_string()), + (4, "name_4".to_string()), + ], + "id=2 deleted; every survivor keeps its non-nullable value" + ); + + writer.close().await.unwrap(); + } + + /// The storage schema no longer rejects a caller's null, so `put` is the + /// only thing standing between a null and a non-nullable base column. + #[tokio::test] + async fn test_put_rejects_null_in_non_nullable_base_column() { + let (store, base_path, base_uri, _temp) = create_local_store().await; + let schema = create_strict_pk_test_schema(); + let writer = ShardWriter::open( + store, + base_path, + base_uri, + ShardWriterConfig { + shard_id: Uuid::new_v4(), + ..Default::default() + }, + schema.clone(), + vec![], + ) + .await + .unwrap(); + + let error = writer.put(vec![null_name_batch()]).await.unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected InvalidInput, got {error:?}" + ); + assert!( + error.to_string().contains("base table schema"), + "error should point at the schema contract: {error}" + ); + + writer.close().await.unwrap(); + } + + /// WAL-only mode validates too — it has no memtable, so before this gate + /// nothing checked its input at all. + #[tokio::test] + async fn test_wal_only_put_rejects_null_in_non_nullable_base_column() { + let (store, base_path, base_uri, _temp) = create_local_store().await; + let schema = create_strict_pk_test_schema(); + let writer = ShardWriter::open( + store, + base_path, + base_uri, + wal_only_config(Uuid::new_v4()), + schema.clone(), + vec![], + ) + .await + .unwrap(); + + let error = writer.put(vec![null_name_batch()]).await.unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected InvalidInput, got {error:?}" + ); + + writer.close().await.unwrap(); + } + + /// A caller-shaped batch that declares `name` nullable and carries a null — + /// legal Arrow, illegal against a base table that declares it non-nullable. + fn null_name_batch() -> RecordBatch { + 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![0, 1])), + Arc::new(StringArray::from(vec![Some("a"), None])), + ], + ) + .unwrap() + } + /// `delete_no_wait` lands the tombstone in the in-memory tier (visible at /// the batch-store level the instant it returns) and hands back the /// durability watcher *without* awaiting it. Index-driven LSM read