From 9a9a2367bd1e5a00c670899cfb75c22c3546384b Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 6 Aug 2026 10:32:28 -0500 Subject: [PATCH 1/3] feat(mem_wal): support delete against non-nullable base columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A tombstone carries the primary key and null in every other column, so `delete` previously required every non-PK column to be nullable in the base table — pushing a storage-engine detail into the user's schema. Split the shard's schema in two. The *logical* schema is the base table's, exactly as the caller declared it; the *storage* schema widens every non-PK top-level field to nullable and is what the memtable, WAL entries, and SSTables physically carry. This mirrors the logical/physical split `SchemaAdapter` already applies to JSON and view types. Widening is top-level only: Arrow validates nullability just at the top level of a `RecordBatch`, so a null `FixedSizeList` or `Struct` needs no change below the top and a vector column's item field gains no validity layer. Primary keys are never widened, so `build_tombstone_batch` still rejects a null, mistyped, or missing key with no extra check. The contract is enforced at two boundaries: * Ingress — `put` validates caller input against the logical schema before the WAL append. This is now the only gate: append and `merge_insert` both compare schemas with `NullabilityComparison::Ignore`, and the encoder derives validity from the array rather than the field, so a null that got past here would reach the base table silently. Validating pre-append also keeps a rejected batch from wedging replay. WAL-only mode is covered too; it previously validated nothing. * Egress — the scan path narrows back to the logical schema once tombstone rows have been filtered out. `project_to_canonical` now actually emits its `target_schema` (DataFusion derives `ProjectionExec` nullability from its expressions, so it could not before), via a new `SchemaRelabelExec`. The same node widens arms so they agree before `UnionExec`/`CoalesceFirstExec`, which require exact schema equality. Narrowing doubles as the assertion that no tombstone escaped its filter. `ensure_tombstone_column` now always re-labels rather than returning a batch that already has the column unchanged, so an entry written under an older storage schema replays into the current one. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/mem_wal.rs | 138 ++++++++ .../lance/src/dataset/mem_wal/scanner/exec.rs | 3 + .../mem_wal/scanner/exec/schema_relabel.rs | 245 +++++++++++++ .../dataset/mem_wal/scanner/point_lookup.rs | 42 ++- .../src/dataset/mem_wal/scanner/projection.rs | 28 +- rust/lance/src/dataset/mem_wal/write.rs | 333 ++++++++++++++++-- 6 files changed, 740 insertions(+), 49 deletions(-) create mode 100644 rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 784ee1fa76c..63892b820bd 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -74,6 +74,52 @@ pub fn tombstone_field() -> ArrowField { ArrowField::new(TOMBSTONE, DataType::Boolean, false) } +/// Derive a shard's **storage schema** from its **logical schema** by making +/// every top-level field nullable except the primary key and `_tombstone`. +/// +/// A shard carries two schemas. The *logical* schema is the base table's, as +/// the caller declared it: it is the contract [`write::ShardWriter::put`] +/// validates input against, and the schema the scan path narrows back to at its +/// output boundary. The *storage* schema is what the memtable, WAL entries, and +/// SSTables physically hold — it exists because a tombstone carries the primary +/// key and null in every other column, so the storage tier must permit a null +/// wherever the base table does not. This mirrors the logical/physical split +/// `SchemaAdapter` already applies to JSON and view types in +/// `crate::dataset::utils`. +/// +/// Top-level only: Arrow validates nullability just at the top level of a +/// `RecordBatch`, so a null `FixedSizeList` or `Struct` needs no change to its +/// child fields — they stay exactly as the caller declared them, and a vector +/// column's item field gains no validity layer. +/// +/// The primary key is excluded because [`lance_core::datatypes::Schema`] +/// requires PK fields and their ancestors to be non-nullable, and a tombstone +/// always carries a real key. `_tombstone` is excluded because the write path +/// always populates it. +/// +/// Idempotent: relaxing an already-relaxed schema is a no-op. +pub fn relax_non_pk_nullability(logical: &ArrowSchema, pk_columns: &[String]) -> Arc { + let fields: Vec = logical + .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.metadata().clone(), + )) +} + /// Extend a base schema with the trailing `_tombstone` column to form the /// mem_wal memtable/generation schema. /// @@ -104,3 +150,95 @@ 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 of a RecordBatch, so + // a null struct/vector needs no change below the top — 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.clone(), 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..6824646995d 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 storage/logical 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..931c26e07c1 --- /dev/null +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/schema_relabel.rs @@ -0,0 +1,245 @@ +// 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, keeping the column arrays +/// untouched. +/// +/// A shard's storage tier widens non-PK columns to nullable so a tombstone can +/// null them (see `relax_non_pk_nullability`), so an LSM plan's arms can +/// disagree with the base-table arm on nullability alone. DataFusion derives a +/// `ProjectionExec`'s output nullability from its expressions rather than from +/// the schema the planner intended, so it cannot pin that down on its own — +/// hence this node. +/// +/// It is used in both directions: +/// +/// - **Widening** (non-nullable → nullable), to make every arm agree before a +/// `UnionExec` or `CoalesceFirstExec`. Always succeeds. +/// - **Narrowing** (nullable → non-nullable), once at the scan's output +/// boundary, to restore the base table's declared schema. This is also the +/// assertion that no tombstone row escaped its filter: `RecordBatch::try_new` +/// rejects a null in a column the target declares non-nullable, so a leak +/// surfaces as an error instead of a row of nulls 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`. + /// + /// The caller is responsible for `schema` being column-compatible with the + /// input (same count, same order, same data types); only field names, + /// nullability, and metadata may differ. A mismatch surfaces per batch at + /// execution time rather than at plan time, because the arrays are what + /// `RecordBatch::try_new` actually validates. + 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(); + // An empty batch carries no arrays to re-label and `try_new` + // cannot infer its row count, so hand back an empty batch with + // the target schema directly. + 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() { + // The post-tombstone-filter case: the storage tier declared `name` + // nullable, 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 row that escaped its filter must not reach the caller as + // a row of nulls — the narrowing relabel is what catches it. + 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..3acd539088b 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,21 @@ 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, widened to the shard's storage nullability, +/// plus a trailing non-nullable `_tombstone` Boolean. +/// +/// Widened because tombstone rows are still in flight at this point — they are +/// only dropped after `CoalesceFirstExec` picks a winner (see +/// [`filter_tombstones_after_coalesce`]), and a tombstone carries null in every +/// non-PK column. Narrowing back to the base table's nullability happens on the +/// far side of that filter. +/// +/// `_tombstone` itself 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 relaxed = relax_non_pk_nullability(canonical, pk_columns); + let mut fields: Vec> = relaxed.fields().iter().cloned().collect(); fields.push(Arc::new(Field::new(TOMBSTONE, DataType::Boolean, false))); Arc::new(Schema::new(fields)) } @@ -772,9 +782,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 +809,12 @@ 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)) + })?); + // Every arm must land on exactly `carry`, or `CoalesceFirstExec` panics on + // the nullability difference between the base arm and the WAL arms. + 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..79eb942bffd 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,33 @@ pub fn null_columns( Ok(Arc::new(projection_exec)) } +/// Force `plan` to report exactly `target_schema`, re-labeling its batches when +/// the plan's own schema differs. +/// +/// DataFusion derives a `ProjectionExec`'s field nullability from its +/// expressions (`Column::nullable(input_schema)`), not from the schema the +/// planner asked for, so a projection alone cannot pin this down. It matters +/// because a shard's storage tier widens non-PK columns to nullable while the +/// base-table arm keeps the base table's nullability, and `CoalesceFirstExec` +/// and `concat_batches` both require exact schema equality. +/// +/// A no-op when the schemas already agree. +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. +/// +/// The result reports `target_schema` exactly, including nullability — see +/// [`force_schema`]. pub fn project_to_canonical( plan: Arc, target_schema: &SchemaRef, @@ -222,7 +248,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..638b29481c9 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; @@ -1045,23 +1045,24 @@ 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 shard's storage schema, injecting `_tombstone = +/// false` for every row when the column is absent. /// /// 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. +/// `_tombstone` keeps its values and is re-labeled rather than passed through, +/// so an entry written under an older, narrower storage schema still replays +/// into the current one. fn ensure_tombstone_column( batch: RecordBatch, target_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]))); + if batch.schema().column_with_name(TOMBSTONE).is_none() { + columns.push(Arc::new(BooleanArray::from(vec![false; n]))); + } RecordBatch::try_new(target_schema.clone(), columns).map_err(|e| { Error::invalid_input(format!( "failed to inject _tombstone column (does the batch match the base schema?): {}", @@ -1072,11 +1073,13 @@ fn ensure_tombstone_column( /// 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. +/// column in the storage schema is 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. +/// `target_schema` is the storage schema, whose non-PK columns are nullable +/// regardless of what the base table declares — that is what lets a tombstone +/// exist for a base table with non-nullable columns. Primary keys stay +/// non-nullable there, so the `RecordBatch` validation below still rejects a +/// null, mistyped, or missing key. fn build_tombstone_batch( keys: &RecordBatch, target_schema: &Arc, @@ -1102,7 +1105,7 @@ fn build_tombstone_batch( } RecordBatch::try_new(target_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 +1504,13 @@ pub struct ShardWriter { manifest_store: Arc, stats: SharedWriteStats, mode: WriterMode, + /// The base table's schema exactly as the caller passed it — no + /// `_tombstone`, nullability untouched. The shard's *logical* half of the + /// schema pair: it is what caller input is held to and what the scan path + /// narrows back to, while the storage schema + /// ([`relax_non_pk_nullability`]) is what the memtable, WAL, and SSTables + /// physically carry. See [`Self::validate_against_logical_schema`]. + logical_schema: Arc, } impl ShardWriter { @@ -1539,10 +1549,14 @@ 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); + // Callers pass the base table's schema. It becomes the shard's *logical* + // schema — the contract every batch handed to `put` is validated against + // — while the *storage* schema the memtable, WAL, and SSTables actually + // carry is derived from it below, once the primary key is known. lance + // owns the `_tombstone` column 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 +1574,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 +1583,19 @@ 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 after the primary key is known: a tombstone nulls every + // non-PK column, so the storage tier has to allow a null where the + // base table does not. `unenforced_primary_key` above ran against the + // unrelaxed 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 +1656,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 +1697,7 @@ impl ShardWriter { manifest_store, stats, mode, + logical_schema, }) } @@ -1954,6 +1979,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 +1987,9 @@ 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. + // Inject `_tombstone = false` and re-label to the storage + // schema; callers only ever pass base-shaped batches and never + // name the column. let batches = batches .into_iter() .map(|b| ensure_tombstone_column(b, &writer_state.schema)) @@ -1992,9 +2018,10 @@ 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. The base table's own nullability is not + /// a constraint here: tombstones live in the shard's storage schema, whose + /// non-PK columns are widened to nullable for exactly this reason, so a + /// delete works against a base table with non-nullable columns. /// /// ``` /// # use lance::Result; @@ -2081,6 +2108,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,7 +2116,7 @@ impl ShardWriter { writer_state, backpressure, } => { - // Inject `_tombstone = false` to match the extended memtable + // Inject `_tombstone = false` and re-label to the storage // schema, mirroring `put`. let batches = batches .into_iter() @@ -2103,6 +2131,33 @@ impl ShardWriter { } } + /// Reject caller input that violates the shard's logical schema: wrong + /// column count, wrong types, or a null in a column the base table declares + /// non-nullable. + /// + /// This is the *only* gate on that contract. The storage schema widens + /// every non-PK column so a tombstone can null it, so it no longer rejects + /// a caller's null, and nothing downstream would either: both append and + /// `merge_insert` compare schemas with `NullabilityComparison::Ignore`, and + /// the encoder derives validity from the array rather than from the field. + /// A null that gets past here reaches the base table silently. + /// + /// Runs before the WAL append, not after: a batch that is appended and only + /// then rejected would fail identically on every subsequent 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 +3734,17 @@ mod tests { ])) } + /// [`create_pk_test_schema`] with a **non-nullable** `name`: the shape that + /// used to make `delete` fail, since a tombstone has to null that column. + 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( @@ -3756,14 +3822,58 @@ mod tests { } #[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 target = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); + + let batch = build_tombstone_batch(&id_only_keys(&[1]), &target, &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 target = 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, &target, &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 +3942,163 @@ mod tests { writer.close().await.unwrap(); } + /// A base table whose non-PK columns are all non-nullable still supports + /// delete: tombstones live in the widened storage schema, and the surviving + /// rows come back through the narrowing egress relabel with their values + /// 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 { + // The scan must hand back the base table's own nullability, not the + // widened storage schema. + 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 From 5c7449c5f19a3c303f9469d9a9384b361b879c02 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Fri, 7 Aug 2026 09:46:01 -0500 Subject: [PATCH 2/3] refactor(mem_wal): tighten schema comments, unify logical/storage terms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cut the comments added by the logical/storage schema split roughly in half — keeping the why, dropping the restatement — and settle on one vocabulary for the pair. `logical schema` is the base table's, as the caller declared it; `storage schema` is the widened one the memtable, WAL, and SSTables carry. Renames follow: `relax_non_pk_nullability(logical_schema, ..)`, and the `target_schema` parameters of `ensure_tombstone_column` / `build_tombstone_batch` (both always receive the storage schema) plus their test locals. User-facing error text keeps "base table schema", which callers recognize. Two comments were stale rather than merely wordy: WAL replay no longer passes a `_tombstone`-carrying batch through unchanged, and `schema_with_tombstone` now produces the intermediate that gets widened, not the memtable schema itself. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/mem_wal.rs | 57 +++---- .../lance/src/dataset/mem_wal/scanner/exec.rs | 2 +- .../mem_wal/scanner/exec/schema_relabel.rs | 51 +++--- .../dataset/mem_wal/scanner/point_lookup.rs | 22 ++- .../src/dataset/mem_wal/scanner/projection.rs | 19 +-- rust/lance/src/dataset/mem_wal/write.rs | 152 ++++++++---------- 6 files changed, 130 insertions(+), 173 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 63892b820bd..1214c2017ee 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,32 +74,26 @@ pub fn tombstone_field() -> ArrowField { ArrowField::new(TOMBSTONE, DataType::Boolean, false) } -/// Derive a shard's **storage schema** from its **logical schema** by making -/// every top-level field nullable except the primary key and `_tombstone`. +/// 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 shard carries two schemas. The *logical* schema is the base table's, as -/// the caller declared it: it is the contract [`write::ShardWriter::put`] -/// validates input against, and the schema the scan path narrows back to at its -/// output boundary. The *storage* schema is what the memtable, WAL entries, and -/// SSTables physically hold — it exists because a tombstone carries the primary -/// key and null in every other column, so the storage tier must permit a null -/// wherever the base table does not. This mirrors the logical/physical split -/// `SchemaAdapter` already applies to JSON and view types in -/// `crate::dataset::utils`. +/// 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 just at the top level of a -/// `RecordBatch`, so a null `FixedSizeList` or `Struct` needs no change to its -/// child fields — they stay exactly as the caller declared them, and a vector -/// column's item field gains no validity layer. -/// -/// The primary key is excluded because [`lance_core::datatypes::Schema`] -/// requires PK fields and their ancestors to be non-nullable, and a tombstone -/// always carries a real key. `_tombstone` is excluded because the write path -/// always populates it. -/// -/// Idempotent: relaxing an already-relaxed schema is a no-op. -pub fn relax_non_pk_nullability(logical: &ArrowSchema, pk_columns: &[String]) -> Arc { - let fields: Vec = logical +/// 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| { @@ -116,12 +110,12 @@ pub fn relax_non_pk_nullability(logical: &ArrowSchema, pk_columns: &[String]) -> .collect(); Arc::new(ArrowSchema::new_with_metadata( fields, - logical.metadata().clone(), + logical_schema.metadata().clone(), )) } -/// Extend a base schema with the trailing `_tombstone` column to form the -/// mem_wal memtable/generation schema. +/// 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 @@ -184,8 +178,7 @@ mod tests { #[test] fn relax_leaves_nested_fields_exactly_as_declared() { - // Arrow validates nullability only at the top level of a RecordBatch, so - // a null struct/vector needs no change below the top — and a vector + // 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); diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/scanner/exec.rs index 6824646995d..9c47c893d8d 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec.rs @@ -10,7 +10,7 @@ //! - [`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 storage/logical nullability boundary) +//! - [`SchemaRelabelExec`]: Re-labels batches to an exact schema (the logical/storage nullability boundary) mod bloom_guard; mod coalesce_first; 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 931c26e07c1..33131c95d5e 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 @@ -19,25 +19,20 @@ use datafusion::physical_plan::{ }; use futures::{Stream, StreamExt}; -/// Re-labels every batch to an exact target schema, keeping the column arrays +/// Re-labels every batch to an exact target schema, leaving the arrays /// untouched. /// -/// A shard's storage tier widens non-PK columns to nullable so a tombstone can -/// null them (see `relax_non_pk_nullability`), so an LSM plan's arms can -/// disagree with the base-table arm on nullability alone. DataFusion derives a -/// `ProjectionExec`'s output nullability from its expressions rather than from -/// the schema the planner intended, so it cannot pin that down on its own — -/// hence this node. +/// 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. /// -/// It is used in both directions: -/// -/// - **Widening** (non-nullable → nullable), to make every arm agree before a -/// `UnionExec` or `CoalesceFirstExec`. Always succeeds. -/// - **Narrowing** (nullable → non-nullable), once at the scan's output -/// boundary, to restore the base table's declared schema. This is also the -/// assertion that no tombstone row escaped its filter: `RecordBatch::try_new` -/// rejects a null in a column the target declares non-nullable, so a leak -/// surfaces as an error instead of a row of nulls reaching the caller. +/// 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, @@ -46,13 +41,10 @@ pub struct SchemaRelabelExec { } impl SchemaRelabelExec { - /// Wrap `input` so its batches are re-labeled to `schema`. - /// - /// The caller is responsible for `schema` being column-compatible with the - /// input (same count, same order, same data types); only field names, - /// nullability, and metadata may differ. A mismatch surfaces per batch at - /// execution time rather than at plan time, because the arrays are what - /// `RecordBatch::try_new` actually validates. + /// 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()), @@ -136,9 +128,8 @@ impl Stream for SchemaRelabelStream { match self.input.poll_next_unpin(cx) { Poll::Ready(Some(Ok(batch))) => { let schema = self.schema.clone(); - // An empty batch carries no arrays to re-label and `try_new` - // cannot infer its row count, so hand back an empty batch with - // the target schema directly. + // 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 { @@ -209,8 +200,8 @@ mod tests { #[tokio::test] async fn narrowing_succeeds_when_no_nulls_remain() { - // The post-tombstone-filter case: the storage tier declared `name` - // nullable, but every surviving row has a value. + // 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))); @@ -221,8 +212,8 @@ mod tests { #[tokio::test] async fn narrowing_rejects_a_surviving_null() { - // A tombstone row that escaped its filter must not reach the caller as - // a row of nulls — the narrowing relabel is what catches it. + // 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))); 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 3acd539088b..fd9daa73756 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -756,21 +756,18 @@ fn cols_with_tombstone(cols: &[String], present: bool) -> Vec { out } -/// Carry schema = canonical output, widened to the shard's storage nullability, -/// plus a trailing non-nullable `_tombstone` Boolean. +/// Carry schema = canonical output at the storage schema's nullability, plus a +/// trailing non-nullable `_tombstone` Boolean. /// -/// Widened because tombstone rows are still in flight at this point — they are -/// only dropped after `CoalesceFirstExec` picks a winner (see -/// [`filter_tombstones_after_coalesce`]), and a tombstone carries null in every -/// non-PK column. Narrowing back to the base table's nullability happens on the -/// far side of that filter. -/// -/// `_tombstone` itself stays non-nullable so the base arm's synthesized +/// 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 relaxed = relax_non_pk_nullability(canonical, pk_columns); - let mut fields: Vec> = relaxed.fields().iter().cloned().collect(); + 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)) } @@ -812,8 +809,7 @@ fn project_to_carry( let projected = Arc::new(ProjectionExec::try_new(project_exprs, plan).map_err(|e| { lance_core::Error::internal(format!("Failed to build carry ProjectionExec: {}", e)) })?); - // Every arm must land on exactly `carry`, or `CoalesceFirstExec` panics on - // the nullability difference between the base arm and the WAL arms. + // `CoalesceFirstExec` panics unless every arm lands on exactly `carry`. Ok(force_schema(projected, &carry)) } diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs index 79eb942bffd..48c0d655496 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs @@ -191,17 +191,13 @@ pub fn null_columns( Ok(Arc::new(projection_exec)) } -/// Force `plan` to report exactly `target_schema`, re-labeling its batches when -/// the plan's own schema differs. +/// Force `plan` to report exactly `target_schema`; a no-op when they agree. /// -/// DataFusion derives a `ProjectionExec`'s field nullability from its -/// expressions (`Column::nullable(input_schema)`), not from the schema the -/// planner asked for, so a projection alone cannot pin this down. It matters -/// because a shard's storage tier widens non-PK columns to nullable while the -/// base-table arm keeps the base table's nullability, and `CoalesceFirstExec` -/// and `concat_batches` both require exact schema equality. -/// -/// A no-op when the schemas already 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, @@ -216,8 +212,7 @@ pub(super) fn force_schema( /// forwarded by name; system / `_distance` cols missing from the source /// are NULL-filled. Other missing columns are an internal error. /// -/// The result reports `target_schema` exactly, including nullability — see -/// [`force_schema`]. +/// Reports `target_schema` exactly, nullability included — see [`force_schema`]. pub fn project_to_canonical( plan: Arc, target_schema: &SchemaRef, diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 638b29481c9..2b360b3633a 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -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,49 +1045,45 @@ fn pk_index_columns(pk_columns: &[String], pk_field_ids: &[i32]) -> Vec<(String, .collect() } -/// Re-label `batch` to the shard's storage schema, injecting `_tombstone = -/// false` for every row when the column 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` keeps its values and is re-labeled rather than passed through, -/// so an entry written under an older, narrower storage schema still replays -/// into the current one. +/// 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 { 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]))); } - 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 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 storage schema is null. +/// Build a tombstone batch from a key-only `keys` batch: primary keys carried +/// through, `_tombstone` true, every other column null. /// -/// `target_schema` is the storage schema, whose non-PK columns are nullable -/// regardless of what the base table declares — that is what lets a tombstone -/// exist for a base table with non-nullable columns. Primary keys stay -/// non-nullable there, so the `RecordBatch` validation below still rejects a -/// null, mistyped, or missing key. +/// 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]))); @@ -1103,7 +1099,7 @@ 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 (do the delete keys match the primary key?): {}", e @@ -1504,12 +1500,11 @@ pub struct ShardWriter { manifest_store: Arc, stats: SharedWriteStats, mode: WriterMode, - /// The base table's schema exactly as the caller passed it — no - /// `_tombstone`, nullability untouched. The shard's *logical* half of the - /// schema pair: it is what caller input is held to and what the scan path - /// narrows back to, while the storage schema - /// ([`relax_non_pk_nullability`]) is what the memtable, WAL, and SSTables - /// physically carry. See [`Self::validate_against_logical_schema`]. + /// 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, } @@ -1549,12 +1544,10 @@ impl ShardWriter { )); } - // Callers pass the base table's schema. It becomes the shard's *logical* - // schema — the contract every batch handed to `put` is validated against - // — while the *storage* schema the memtable, WAL, and SSTables actually - // carry is derived from it below, once the primary key is known. lance - // owns the `_tombstone` column and appends it here. Idempotent, so a - // reopen that already extended the schema is a no-op. + // 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); @@ -1590,10 +1583,9 @@ impl ShardWriter { &pk_columns, )?; - // Widen only after the primary key is known: a tombstone nulls every - // non-PK column, so the storage tier has to allow a null where the - // base table does not. `unenforced_primary_key` above ran against the - // unrelaxed schema, which is what enforces non-nullable PKs. + // 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 { @@ -1987,9 +1979,8 @@ impl ShardWriter { writer_state, backpressure, } => { - // Inject `_tombstone = false` and re-label to the storage - // 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)) @@ -2018,10 +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. The base table's own nullability is not - /// a constraint here: tombstones live in the shard's storage schema, whose - /// non-PK columns are widened to nullable for exactly this reason, so a - /// delete works against a base table with non-nullable columns. + /// 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; @@ -2116,8 +2105,7 @@ impl ShardWriter { writer_state, backpressure, } => { - // Inject `_tombstone = false` and re-label to the storage - // schema, mirroring `put`. + // Mirrors `put`. let batches = batches .into_iter() .map(|b| ensure_tombstone_column(b, &writer_state.schema)) @@ -2131,20 +2119,18 @@ impl ShardWriter { } } - /// Reject caller input that violates the shard's logical schema: wrong - /// column count, wrong types, or a null in a column the base table declares - /// non-nullable. + /// 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. /// - /// This is the *only* gate on that contract. The storage schema widens - /// every non-PK column so a tombstone can null it, so it no longer rejects - /// a caller's null, and nothing downstream would either: both append and + /// 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 derives validity from the array rather than from the field. - /// A null that gets past here reaches the base table silently. + /// 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, not after: a batch that is appended and only - /// then rejected would fail identically on every subsequent replay, leaving - /// the shard unable to reopen. + /// 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( @@ -3734,8 +3720,8 @@ mod tests { ])) } - /// [`create_pk_test_schema`] with a **non-nullable** `name`: the shape that - /// used to make `delete` fail, since a tombstone has to null that column. + /// [`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() @@ -3760,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() @@ -3774,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") @@ -3808,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", @@ -3818,7 +3804,7 @@ 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] @@ -3830,9 +3816,9 @@ mod tests { Field::new("id", DataType::Int32, false), Field::new("v", DataType::Int32, false), ])); - let target = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); + let storage = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); - let batch = build_tombstone_batch(&id_only_keys(&[1]), &target, &pk).unwrap(); + 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"); @@ -3854,7 +3840,7 @@ mod tests { Field::new("id", DataType::Int32, false), Field::new("v", DataType::Int32, false), ])); - let target = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); + let storage = relax_non_pk_nullability(&schema_with_tombstone(&base), &pk); let keys = RecordBatch::try_new( Arc::new(ArrowSchema::new(vec![Field::new( "id", @@ -3865,7 +3851,7 @@ mod tests { ) .unwrap(); - let error = build_tombstone_batch(&keys, &target, &pk).unwrap_err(); + let error = build_tombstone_batch(&keys, &storage, &pk).unwrap_err(); assert!( matches!(error, Error::InvalidInput { .. }), "expected InvalidInput, got {error:?}" @@ -3942,10 +3928,8 @@ mod tests { writer.close().await.unwrap(); } - /// A base table whose non-PK columns are all non-nullable still supports - /// delete: tombstones live in the widened storage schema, and the surviving - /// rows come back through the narrowing egress relabel with their values - /// intact. + /// 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; @@ -3998,8 +3982,6 @@ mod tests { let mut rows: Vec<(i32, String)> = Vec::new(); for b in &batches { - // The scan must hand back the base table's own nullability, not the - // widened storage schema. assert!( !b.schema().field_with_name("name").unwrap().is_nullable(), "egress must narrow back to the logical schema" From 5d7df36fef1c43dabad8478682db3d32223bdb81 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Fri, 7 Aug 2026 11:04:11 -0500 Subject: [PATCH 3/3] fix(mem_wal): drop redundant clone in nullability test `item` is not used after the `FixedSizeList` is built, so the clone trips `clippy::redundant_clone` and fails the workspace clippy gate. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/mem_wal.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/mem_wal.rs b/rust/lance/src/dataset/mem_wal.rs index 1214c2017ee..40e7dbde46d 100644 --- a/rust/lance/src/dataset/mem_wal.rs +++ b/rust/lance/src/dataset/mem_wal.rs @@ -184,7 +184,7 @@ mod tests { 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.clone(), 4), false), + ArrowField::new("vector", DataType::FixedSizeList(item, 4), false), ArrowField::new("s", DataType::Struct(Fields::from(vec![child])), false), ]);