From c2a647189d0600e6950e9c6f418bb86b6bd3457e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 16:14:22 +0800 Subject: [PATCH] feat: invalidate generated columns on native append --- .../plan_generated_column_invalidation.rs | 12 +- rust/lancedb/src/table.rs | 34 +- .../append_generated_column_invalidation.rs | 150 ++++ ..._generated_column_invalidation_contract.rs | 792 ++++++++++++++++++ rust/lancedb/src/table/datafusion/insert.rs | 27 +- 5 files changed, 999 insertions(+), 16 deletions(-) create mode 100644 rust/lancedb/src/table/append_generated_column_invalidation.rs create mode 100644 rust/lancedb/src/table/append_generated_column_invalidation_contract.rs diff --git a/rust/lancedb/src/function/plan_generated_column_invalidation.rs b/rust/lancedb/src/function/plan_generated_column_invalidation.rs index 2ccdb965d..7c9c2b60e 100644 --- a/rust/lancedb/src/function/plan_generated_column_invalidation.rs +++ b/rust/lancedb/src/function/plan_generated_column_invalidation.rs @@ -5,12 +5,8 @@ //! //! Plans column-wide dependency-epoch advances from a binding snapshot and a //! mutation impact. This module does not mutate tables, write metadata, or -//! execute append/update/delete/merge paths. -//! -//! The planner surface is intentionally unwired from Native mutation paths in -//! this slice; suppress unused warnings outside tests until that consumer lands. - -#![cfg_attr(not(test), allow(dead_code))] +//! execute append/update/delete/merge paths. Native append consumes the plan +//! through the B4b runtime wiring. use std::collections::BTreeSet; @@ -23,6 +19,10 @@ pub enum GeneratedColumnMutationImpact { /// Append or delete: whole-column coverage / row membership changed. RowSetChanged, /// Update of the listed stable field IDs (direct and transitive dependents). + /// + /// Reserved for future update/delete/merge invalidation consumers; Native + /// append (B4b) only constructs [`Self::RowSetChanged`]. + #[cfg_attr(not(test), allow(dead_code))] UpdatedFields(BTreeSet), } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 7f30b84e2..a79f44285 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -70,6 +70,7 @@ use self::merge::MergeInsertBuilder; pub mod add_columns; mod add_data; +mod append_generated_column_invalidation; pub mod branch_merge; pub mod checkpoint; mod create_index; @@ -85,6 +86,8 @@ pub mod schema_evolution; pub mod update; pub mod write_progress; +#[cfg(test)] +mod append_generated_column_invalidation_contract; #[cfg(test)] mod schema_metadata_updates_dependency_contract; @@ -3221,13 +3224,14 @@ impl BaseTable for NativeTable { } async fn add(&self, mut add: AddDataBuilder) -> Result { - let table_def = self.table_definition().await?; - self.dataset.ensure_mutable()?; let ds_wrapper = self.dataset.clone(); + // One exact dataset supplies table definition, schema, binding snapshot, + // planning basis, and the InsertExec transaction basis. let ds = self.dataset.get().await?; - let table_schema = Schema::from(&ds.schema().clone()); + let table_schema = Schema::from(ds.schema()); + let table_def = TableDefinition::try_from_rich_schema(Arc::new(table_schema.clone()))?; let num_partitions = if let Some(parallelism) = add.write_parallelism { parallelism @@ -3263,6 +3267,14 @@ impl BaseTable for NativeTable { ..Default::default() }); + // Plan after in-memory preprocessing, before any InsertExec file write. + // Canonical overwrite is PreprocessingOutput.overwrite, not final lance_params.mode. + let schema_metadata_updates = + append_generated_column_invalidation::plan_native_append_generated_column_invalidation( + ds.as_ref(), + output.overwrite, + )?; + // Repartition for write parallelism if beneficial. let plan = if num_partitions > 1 { Arc::new( @@ -3281,6 +3293,7 @@ impl BaseTable for NativeTable { plan, lance_params, output.tracker.clone(), + schema_metadata_updates, )); let tracker_for_tasks = output.tracker.clone(); @@ -3773,13 +3786,22 @@ impl BaseTable for NativeTable { input: Arc, write_params: WriteParams, ) -> Result> { - let ds = self.dataset.get().await?; - let dataset = Arc::new((*ds).clone()); - Ok(Arc::new(datafusion::insert::InsertExec::new( + // One exact dataset supplies planning basis and the InsertExec basis. + let dataset = self.dataset.get().await?; + let is_overwrite = matches!(write_params.mode, WriteMode::Overwrite); + // Reject generated-table overwrite before returning an execution plan. + let schema_metadata_updates = + append_generated_column_invalidation::plan_native_append_generated_column_invalidation( + dataset.as_ref(), + is_overwrite, + )?; + Ok(Arc::new(datafusion::insert::InsertExec::new_with_tracker( self.dataset.clone(), dataset, input, write_params, + None, + schema_metadata_updates, ))) } } diff --git a/rust/lancedb/src/table/append_generated_column_invalidation.rs b/rust/lancedb/src/table/append_generated_column_invalidation.rs new file mode 100644 index 000000000..af69a35cb --- /dev/null +++ b/rust/lancedb/src/table/append_generated_column_invalidation.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Crate-private Native append wiring for generated-column invalidation (B4b). +//! +//! Converts the B4a pure planner into one Lance field-metadata patch for Native +//! append commits. Planning is strict-decode/validate; overwrite of a table with +//! any generated-column definition fails closed as [`Error::NotSupported`]. + +use std::collections::HashMap; + +use lance::Dataset; +use lance::dataset::transaction::{SchemaMetadataUpdates, UpdateMap, UpdateMapEntry}; + +use crate::Result; +use crate::error::Error; +use crate::function::GENERATED_COLUMN_METADATA_KEY; +use crate::function::plan_generated_column_invalidation::{ + GeneratedColumnMutationImpact, PlannedGeneratedColumnMetadataUpdate, + plan_generated_column_invalidation, +}; + +use super::generated_column_binding_snapshot_from_dataset; + +/// Plan Native append invalidation against one exact dataset snapshot. +/// +/// Strict-decodes and validates every present generated-column metadata value +/// through the B4a planner. When `is_overwrite` is true and any generated column +/// is present, returns [`Error::NotSupported`] before mutation. Otherwise returns +/// `Some(patch)` when at least one generated column would be invalidated, or +/// `None` when the table has no generated columns. +pub(super) fn plan_native_append_generated_column_invalidation( + dataset: &Dataset, + is_overwrite: bool, +) -> Result> { + let snapshot = generated_column_binding_snapshot_from_dataset(dataset)?; + let plan = plan_generated_column_invalidation( + &snapshot, + &GeneratedColumnMutationImpact::RowSetChanged, + )?; + if plan.is_empty() { + return Ok(None); + } + if is_overwrite { + return Err(Error::NotSupported { + message: "Overwrite is not supported on tables with generated columns".to_string(), + }); + } + Ok(Some(planned_invalidation_to_schema_metadata_updates(plan))) +} + +/// Convert planner replacements into one non-empty Lance field-metadata patch. +/// +/// Each entry is keyed by stable output field ID, uses `replace: false`, and +/// replaces only [`GENERATED_COLUMN_METADATA_KEY`]. +fn planned_invalidation_to_schema_metadata_updates( + plan: Vec, +) -> SchemaMetadataUpdates { + SchemaMetadataUpdates { + schema_metadata_updates: None, + field_metadata_updates: plan + .into_iter() + .map(|update| { + ( + update.output_field_id(), + UpdateMap { + update_entries: vec![UpdateMapEntry { + key: GENERATED_COLUMN_METADATA_KEY.to_string(), + value: Some(update.metadata_json().to_string()), + }], + replace: false, + }, + ) + }) + .collect::>(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Construct a planned update through the public accessors by planning a + /// minimal in-memory snapshot, then assert the Lance patch shape. + #[test] + fn planned_replacements_become_non_replace_field_patch() { + use crate::function::{ + Function, FunctionArgument, FunctionCall, FunctionId, FunctionOutput, + FunctionParameter, FunctionSignature, GeneratedColumnBindingSnapshot, + GeneratedColumnDefinition, + }; + use arrow_array::{ArrayRef, StringArray}; + use arrow_schema::{DataType, Field}; + use std::sync::Arc; + + let field_id = 11; + let function = Function::new( + FunctionId::try_new("fn.exact.b4b.helper.patch").unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ); + let call = FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("x")])) as ArrayRef + ) + .unwrap(), + )], + ) + .unwrap(); + let definition = GeneratedColumnDefinition::try_new(field_id, call, 3, 3).unwrap(); + let json = definition.to_metadata_json().unwrap(); + let snap = GeneratedColumnBindingSnapshot::try_new( + 1, + vec![Arc::new( + Field::new("gen_out", DataType::Int32, true) + .with_metadata([(GENERATED_COLUMN_METADATA_KEY.to_string(), json)].into()), + )], + vec![field_id], + ) + .unwrap(); + let plan = plan_generated_column_invalidation( + &snap, + &GeneratedColumnMutationImpact::RowSetChanged, + ) + .unwrap(); + let patch = planned_invalidation_to_schema_metadata_updates(plan); + assert!(!patch.is_empty()); + assert!(patch.schema_metadata_updates.is_none()); + let map = patch + .field_metadata_updates + .get(&field_id) + .expect("stable field id must be present"); + assert!(!map.replace); + assert_eq!(map.update_entries.len(), 1); + assert_eq!(map.update_entries[0].key, GENERATED_COLUMN_METADATA_KEY); + let decoded = GeneratedColumnDefinition::from_metadata_json( + map.update_entries[0].value.as_deref().unwrap(), + field_id, + ) + .unwrap(); + assert_eq!(decoded.dependency_epoch(), 4); + assert_eq!(decoded.materialized_epoch(), 3); + } +} diff --git a/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs b/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs new file mode 100644 index 000000000..323858f3e --- /dev/null +++ b/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs @@ -0,0 +1,792 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! RED runtime contract tests for Native append invalidation (B4b). +//! +//! These tests pin Native Table API and DataFusion SQL insert behavior for +//! generated-column dependency-epoch invalidation. They use real local Native +//! tables and existing public/internal APIs; Lance commits and query guards are +//! not mocked. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::{Array, Int32Array, RecordBatch, RecordBatchIterator, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::prelude::SessionContext; +use futures::TryStreamExt; +use lance::dataset::{WriteMode, WriteParams}; +use tempfile::TempDir; + +use crate::connection::ConnectBuilder; +use crate::error::{Error, FunctionErrorCode}; +use crate::function::{ + Function, FunctionArgument, FunctionCall, FunctionId, FunctionOutput, FunctionParameter, + FunctionSignature, GENERATED_COLUMN_METADATA_KEY, GeneratedColumnDefinition, + GeneratedColumnStatus, +}; +use crate::query::{ExecutableQuery, QueryBase, Select}; +use crate::table::datafusion::BaseTableAdapter; +use crate::table::schema_evolution::FieldMetadataUpdate; +use crate::table::{AddDataMode, Table, WriteOptions}; + +const GEN_OUT: &str = "gen_out"; +const ORDINARY: &str = "ordinary"; +const INITIAL_DEPENDENCY_EPOCH: u64 = 3; +const INITIAL_MATERIALIZED_EPOCH: u64 = 3; +const MALFORMED_MARKER: &str = "SENSITIVE_B4B_APPEND_METADATA_MARKER_9f2c_a81d"; + +struct Fixture { + _tmp: TempDir, + table: Table, + table_name: String, + uri: String, +} + +fn literal_only_function() -> Function { + Function::new( + FunctionId::try_new("fn.exact.b4b.append.literal").unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ) +} + +fn literal_only_definition( + output_field_id: i32, + dependency_epoch: u64, + materialized_epoch: u64, +) -> GeneratedColumnDefinition { + let function = literal_only_function(); + let call = FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("literal-only")])) as arrow_array::ArrayRef + ) + .unwrap(), + )], + ) + .unwrap(); + assert!( + call.arguments() + .iter() + .all(|(_, argument)| argument.field_id().is_none()), + "fixture must be literal-only so row-set coverage, not field dependency, drives invalidation" + ); + GeneratedColumnDefinition::try_new(output_field_id, call, dependency_epoch, materialized_epoch) + .unwrap() +} + +async fn create_table_with_complete_literal_generated(name: &str) -> Fixture { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_str().unwrap().to_string(); + let conn = ConnectBuilder::new(&uri).execute().await.unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new(GEN_OUT, DataType::Int32, true), + Field::new(ORDINARY, DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("seed")])), + ], + ) + .unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + let field_id = snapshot.field(GEN_OUT).expect(GEN_OUT).field_id(); + let definition = literal_only_definition( + field_id, + INITIAL_DEPENDENCY_EPOCH, + INITIAL_MATERIALIZED_EPOCH, + ); + let json = definition.to_metadata_json().unwrap(); + table + .update_field_metadata(&[ + FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, json) + ]) + .await + .unwrap(); + + assert_eq!( + table.generated_column_status(GEN_OUT).await.unwrap(), + GeneratedColumnStatus::Complete + ); + let planted = read_generated_definition(&table).await; + assert!( + planted + .function_call() + .arguments() + .iter() + .all(|(_, argument)| argument.field_id().is_none()), + "planted metadata must remain literal-only" + ); + + Fixture { + _tmp: tmp, + table, + table_name: name.to_string(), + uri, + } +} + +async fn create_ordinary_table(name: &str) -> Fixture { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_str().unwrap().to_string(); + let conn = ConnectBuilder::new(&uri).execute().await.unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new(GEN_OUT, DataType::Int32, true), + Field::new(ORDINARY, DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("seed")])), + ], + ) + .unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + Fixture { + _tmp: tmp, + table, + table_name: name.to_string(), + uri, + } +} + +async fn read_generated_definition(table: &Table) -> GeneratedColumnDefinition { + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + snapshot + .field(GEN_OUT) + .expect(GEN_OUT) + .generated_column_definition() + .expect("generated metadata must decode") + .expect("generated metadata must be present") +} + +fn ordinary_rows_batch(values: &[&str]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + ORDINARY, + DataType::Utf8, + true, + )])); + RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from( + values.iter().map(|value| Some(*value)).collect::>(), + ))], + ) + .unwrap() +} + +fn full_rows_batch(gen_values: &[Option], ordinary_values: &[&str]) -> RecordBatch { + assert_eq!(gen_values.len(), ordinary_values.len()); + let schema = Arc::new(Schema::new(vec![ + Field::new(GEN_OUT, DataType::Int32, true), + Field::new(ORDINARY, DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(gen_values.to_vec())), + Arc::new(StringArray::from( + ordinary_values + .iter() + .map(|value| Some(*value)) + .collect::>(), + )), + ], + ) + .unwrap() +} + +async fn ordinary_values(table: &Table) -> HashSet { + let batches = table + .query() + .select(Select::columns(&[ORDINARY])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut values = HashSet::new(); + for batch in batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for index in 0..column.len() { + if !column.is_null(index) { + values.insert(column.value(index).to_string()); + } + } + } + values +} + +fn assert_generated_column_incomplete(err: &Error, label: &str) { + match err { + Error::Function { + code: FunctionErrorCode::GeneratedColumnIncomplete, + .. + } => {} + other => panic!("{label}: expected generated_column_incomplete, got {other:?}"), + } +} + +fn assert_not_supported(err: &Error, label: &str) { + assert!( + matches!(err, Error::NotSupported { .. }), + "{label}: expected NotSupported, got {err:?}" + ); +} + +fn assert_invalid_input_redacted(err: &Error, label: &str) { + assert!( + matches!(err, Error::InvalidInput { .. }), + "{label}: expected InvalidInput, got {err:?}" + ); + let rendered = format!("{err}\n{err:?}"); + assert!( + !rendered.contains(MALFORMED_MARKER), + "{label}: diagnostic echoed unique metadata marker: {rendered}" + ); + assert!( + !rendered.contains(GENERATED_COLUMN_METADATA_KEY), + "{label}: diagnostic echoed metadata wire key: {rendered}" + ); +} + +fn assert_conflict_error(err: &Error, label: &str) { + match err { + Error::Lance { source } => { + assert!( + matches!( + source, + lance::Error::IncompatibleTransaction { .. } + | lance::Error::RetryableCommitConflict { .. } + | lance::Error::CommitConflict { .. } + ), + "{label}: expected Lance commit conflict category, got {source:?}" + ); + } + Error::Function { + code: FunctionErrorCode::StaleOrConflictingInput, + .. + } => {} + other => panic!("{label}: expected conflict error category, got {other:?}"), + } +} + +fn from_datafusion_error(err: datafusion_common::DataFusionError) -> Error { + Error::from(err) +} + +async fn sql_ctx_for(table: &Table, name: &str) -> SessionContext { + let ctx = SessionContext::new(); + let provider = BaseTableAdapter::try_new(table.base_table().clone()) + .await + .unwrap(); + ctx.register_table(name, Arc::new(provider)).unwrap(); + ctx +} + +async fn run_sql(ctx: &SessionContext, sql: &str) -> Result<(), Error> { + match ctx.sql(sql).await { + Err(err) => Err(from_datafusion_error(err)), + Ok(df) => match df.collect().await { + Ok(_) => Ok(()), + Err(err) => Err(from_datafusion_error(err)), + }, + } +} + +#[tokio::test] +async fn nonempty_table_api_append_invalidates_literal_only_generated_column() { + let fixture = create_table_with_complete_literal_generated("b4b_table_append").await; + let before = read_generated_definition(&fixture.table).await; + + fixture + .table + .add(ordinary_rows_batch(&["appended"])) + .execute() + .await + .expect("non-empty Table API append must commit"); + + let values = ordinary_values(&fixture.table).await; + assert!(values.contains("seed")); + assert!(values.contains("appended")); + + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + let after = read_generated_definition(&fixture.table).await; + assert_eq!(after.dependency_epoch(), before.dependency_epoch() + 1); + assert_eq!(after.materialized_epoch(), before.materialized_epoch()); + assert_eq!(after.output_field_id(), before.output_field_id()); + assert_eq!(after.function_call(), before.function_call()); + + let Err(err) = fixture + .table + .query() + .select(Select::columns(&[GEN_OUT])) + .execute() + .await + else { + panic!("incomplete generated column query must fail"); + }; + assert_generated_column_incomplete(&err, "table api append query"); +} + +#[tokio::test] +async fn nonempty_table_api_append_atomic_version_visibility() { + let fixture = create_table_with_complete_literal_generated("b4b_atomic_visibility").await; + let previous_version = fixture.table.version().await.unwrap(); + let previous_rows = ordinary_values(&fixture.table).await; + let previous_definition = read_generated_definition(&fixture.table).await; + assert_eq!( + previous_definition.status(), + GeneratedColumnStatus::Complete + ); + + fixture + .table + .add(ordinary_rows_batch(&["atomic-new"])) + .execute() + .await + .expect("non-empty append must commit"); + let new_version = fixture.table.version().await.unwrap(); + assert_ne!(new_version, previous_version); + + // Exact new version: new rows + incomplete metadata together. + let new_rows = ordinary_values(&fixture.table).await; + assert!(new_rows.contains("atomic-new")); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + let new_definition = read_generated_definition(&fixture.table).await; + assert_eq!( + new_definition.dependency_epoch(), + previous_definition.dependency_epoch() + 1 + ); + + // Immediately previous version: neither new rows nor incomplete metadata. + fixture.table.checkout(previous_version).await.unwrap(); + assert_eq!(ordinary_values(&fixture.table).await, previous_rows); + assert!(!ordinary_values(&fixture.table).await.contains("atomic-new")); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Complete + ); + let checked_out = read_generated_definition(&fixture.table).await; + assert_eq!(checked_out, previous_definition); + fixture + .table + .query() + .select(Select::columns(&[GEN_OUT])) + .execute() + .await + .expect("previous complete version must remain readable"); +} + +#[tokio::test] +async fn empty_table_api_append_leaves_complete_generated_column() { + let fixture = create_table_with_complete_literal_generated("b4b_empty_table_append").await; + let before = read_generated_definition(&fixture.table).await; + let rows_before = ordinary_values(&fixture.table).await; + + fixture + .table + .add(RecordBatch::new_empty(Arc::new(Schema::new(vec![ + Field::new(ORDINARY, DataType::Utf8, true), + ])))) + .execute() + .await + .expect("empty Table API append is a supported path"); + + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Complete + ); + let after = read_generated_definition(&fixture.table).await; + assert_eq!(after, before); +} + +#[tokio::test] +async fn multipartition_table_api_append_advances_dependency_epoch_once() { + let fixture = create_table_with_complete_literal_generated("b4b_multipartition").await; + let before = read_generated_definition(&fixture.table).await; + + fixture + .table + .add(ordinary_rows_batch(&["p0", "p1", "p2", "p3"])) + .write_parallelism(2) + .execute() + .await + .expect("multi-partition append must commit"); + + let values = ordinary_values(&fixture.table).await; + assert!(values.contains("p0")); + assert!(values.contains("p3")); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + let after = read_generated_definition(&fixture.table).await; + assert_eq!( + after.dependency_epoch(), + before.dependency_epoch() + 1, + "multi-partition append must attach one whole-transaction patch" + ); + assert_eq!(after.materialized_epoch(), before.materialized_epoch()); + assert_eq!(after.function_call(), before.function_call()); +} + +#[tokio::test] +async fn table_api_overwrite_rejects_before_mutation_when_generated_column_present() { + let fixture = create_table_with_complete_literal_generated("b4b_table_overwrite").await; + let version_before = fixture.table.version().await.unwrap(); + let rows_before = ordinary_values(&fixture.table).await; + let definition_before = read_generated_definition(&fixture.table).await; + + let err = fixture + .table + .add(full_rows_batch(&[Some(9)], &["overwrite"])) + .mode(AddDataMode::Overwrite) + .execute() + .await + .expect_err("overwrite must reject when any generated column is present"); + assert_not_supported(&err, "table api overwrite"); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert_eq!( + read_generated_definition(&fixture.table).await, + definition_before + ); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Complete + ); +} + +#[tokio::test] +async fn table_api_effective_overwrite_from_add_data_mode_rejects_when_lance_params_append() { + let fixture = + create_table_with_complete_literal_generated("b4b_table_effective_overwrite").await; + let version_before = fixture.table.version().await.unwrap(); + let rows_before = ordinary_values(&fixture.table).await; + let definition_before = read_generated_definition(&fixture.table).await; + assert_eq!(definition_before.status(), GeneratedColumnStatus::Complete); + + let err = fixture + .table + .add(full_rows_batch(&[Some(9)], &["effective-overwrite"])) + .mode(AddDataMode::Overwrite) + .write_options(WriteOptions { + lance_write_params: Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + }) + .execute() + .await + .expect_err( + "AddDataMode::Overwrite must reject generated-table writes even when \ + explicit lance WriteParams.mode is Append", + ); + assert_not_supported(&err, "table api effective overwrite"); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert_eq!( + read_generated_definition(&fixture.table).await, + definition_before + ); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Complete + ); +} + +#[tokio::test] +async fn ordinary_table_api_overwrite_still_supported() { + let fixture = create_ordinary_table("b4b_ordinary_overwrite_control").await; + + fixture + .table + .add(full_rows_batch(&[Some(42)], &["replaced"])) + .mode(AddDataMode::Overwrite) + .execute() + .await + .expect("ordinary tables must keep overwrite support"); + + let values = ordinary_values(&fixture.table).await; + assert_eq!(values, HashSet::from(["replaced".to_string()])); + assert_eq!(fixture.table.count_rows(None).await.unwrap(), 1); +} + +#[tokio::test] +async fn nonempty_sql_insert_invalidates_generated_column() { + let fixture = create_table_with_complete_literal_generated("b4b_sql_insert").await; + let before = read_generated_definition(&fixture.table).await; + let ctx = sql_ctx_for(&fixture.table, &fixture.table_name).await; + + run_sql( + &ctx, + &format!( + "INSERT INTO {} VALUES (CAST(NULL AS INT), 'sql-appended')", + fixture.table_name + ), + ) + .await + .expect("non-empty SQL INSERT must commit"); + + fixture.table.checkout_latest().await.unwrap(); + let values = ordinary_values(&fixture.table).await; + assert!(values.contains("sql-appended")); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Incomplete + ); + let after = read_generated_definition(&fixture.table).await; + assert_eq!(after.dependency_epoch(), before.dependency_epoch() + 1); + assert_eq!(after.materialized_epoch(), before.materialized_epoch()); + assert_eq!(after.function_call(), before.function_call()); + + let Err(err) = fixture + .table + .query() + .select(Select::columns(&[GEN_OUT])) + .execute() + .await + else { + panic!("SQL INSERT invalidation must trip generated query guard"); + }; + assert_generated_column_incomplete(&err, "sql insert query"); +} + +#[tokio::test] +async fn empty_sql_insert_leaves_complete_generated_column() { + let fixture = create_table_with_complete_literal_generated("b4b_empty_sql_insert").await; + let before = read_generated_definition(&fixture.table).await; + let rows_before = ordinary_values(&fixture.table).await; + + let conn = ConnectBuilder::new(&fixture.uri).execute().await.unwrap(); + let source_schema = Arc::new(Schema::new(vec![ + Field::new(GEN_OUT, DataType::Int32, true), + Field::new(ORDINARY, DataType::Utf8, true), + ])); + let empty_reader: Box = + Box::new(RecordBatchIterator::new( + std::iter::empty::>(), + source_schema, + )); + let source = conn + .create_table("empty_source", empty_reader) + .execute() + .await + .unwrap(); + + let ctx = sql_ctx_for(&fixture.table, &fixture.table_name).await; + let source_provider = BaseTableAdapter::try_new(source.base_table().clone()) + .await + .unwrap(); + ctx.register_table("empty_source", Arc::new(source_provider)) + .unwrap(); + + run_sql( + &ctx, + &format!( + "INSERT INTO {} SELECT * FROM empty_source", + fixture.table_name + ), + ) + .await + .expect("empty SQL INSERT is a supported path"); + + fixture.table.checkout_latest().await.unwrap(); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert_eq!( + fixture + .table + .generated_column_status(GEN_OUT) + .await + .unwrap(), + GeneratedColumnStatus::Complete + ); + assert_eq!(read_generated_definition(&fixture.table).await, before); +} + +#[tokio::test] +async fn sql_insert_overwrite_rejects_before_mutation_when_generated_column_present() { + let fixture = create_table_with_complete_literal_generated("b4b_sql_overwrite").await; + let version_before = fixture.table.version().await.unwrap(); + let rows_before = ordinary_values(&fixture.table).await; + let definition_before = read_generated_definition(&fixture.table).await; + let ctx = sql_ctx_for(&fixture.table, &fixture.table_name).await; + + let err = run_sql( + &ctx, + &format!( + "INSERT OVERWRITE INTO {} VALUES (10, 'sql-overwrite')", + fixture.table_name + ), + ) + .await + .expect_err("SQL INSERT OVERWRITE must reject when any generated column is present"); + assert_not_supported(&err, "sql insert overwrite"); + + fixture.table.checkout_latest().await.unwrap(); + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert_eq!( + read_generated_definition(&fixture.table).await, + definition_before + ); +} + +#[tokio::test] +async fn malformed_generated_metadata_rejects_append_before_mutation_and_redacts_marker() { + let fixture = create_ordinary_table("b4b_malformed_preflight").await; + let snapshot = fixture + .table + .generated_column_binding_snapshot() + .await + .unwrap(); + let field_id = snapshot.field(GEN_OUT).expect(GEN_OUT).field_id(); + let raw = format!( + r#"{{"format_version":1,"output_field_id":{field_id},"function_call":{MALFORMED_MARKER},"dependency_epoch":1,"materialized_epoch":1}}"# + ); + assert!(raw.contains(MALFORMED_MARKER)); + fixture + .table + .update_field_metadata(&[ + FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, raw.clone()) + ]) + .await + .unwrap(); + + let version_before = fixture.table.version().await.unwrap(); + let rows_before = ordinary_values(&fixture.table).await; + + let err = fixture + .table + .add(ordinary_rows_batch(&["must-not-land"])) + .execute() + .await + .expect_err("malformed generated metadata must fail closed before append visibility"); + assert_invalid_input_redacted(&err, "malformed append preflight"); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + assert!( + !ordinary_values(&fixture.table) + .await + .contains("must-not-land") + ); +} + +#[tokio::test] +async fn concurrent_same_field_append_one_winner_one_conflict() { + let fixture = create_table_with_complete_literal_generated("b4b_concurrent_append").await; + let conn = ConnectBuilder::new(&fixture.uri) + .read_consistency_interval(Duration::from_secs(3600)) + .execute() + .await + .unwrap(); + let table_a = conn + .open_table(&fixture.table_name) + .execute() + .await + .unwrap(); + let table_b = conn + .open_table(&fixture.table_name) + .execute() + .await + .unwrap(); + let basis_version = table_a.version().await.unwrap(); + assert_eq!(table_b.version().await.unwrap(), basis_version); + + let (result_a, result_b) = tokio::join!( + table_a.add(ordinary_rows_batch(&["winner-a"])).execute(), + table_b.add(ordinary_rows_batch(&["winner-b"])).execute(), + ); + + let outcomes = [result_a, result_b]; + let wins = outcomes.iter().filter(|result| result.is_ok()).count(); + let losses = outcomes.iter().filter(|result| result.is_err()).count(); + assert_eq!(wins, 1, "exactly one same-basis append may publish"); + assert_eq!(losses, 1, "exactly one same-basis append must conflict"); + for result in &outcomes { + if let Err(err) = result { + assert_conflict_error(err, "concurrent same-field append loser"); + } + } + + let fresh = conn + .open_table(&fixture.table_name) + .execute() + .await + .unwrap(); + let values = ordinary_values(&fresh).await; + assert!(values.contains("seed")); + let has_a = values.contains("winner-a"); + let has_b = values.contains("winner-b"); + assert!( + has_a ^ has_b, + "only winner rows may be visible, got {values:?}" + ); + assert_eq!( + fresh.generated_column_status(GEN_OUT).await.unwrap(), + GeneratedColumnStatus::Incomplete + ); + let definition = read_generated_definition(&fresh).await; + assert_eq!(definition.dependency_epoch(), INITIAL_DEPENDENCY_EPOCH + 1); + assert_eq!(definition.materialized_epoch(), INITIAL_MATERIALIZED_EPOCH); +} diff --git a/rust/lancedb/src/table/datafusion/insert.rs b/rust/lancedb/src/table/datafusion/insert.rs index e176c228b..cbf48cf85 100644 --- a/rust/lancedb/src/table/datafusion/insert.rs +++ b/rust/lancedb/src/table/datafusion/insert.rs @@ -19,7 +19,7 @@ use datafusion_physical_plan::{ }; use futures::TryStreamExt; use lance::Dataset; -use lance::dataset::transaction::{Operation, Transaction}; +use lance::dataset::transaction::{Operation, SchemaMetadataUpdates, Transaction}; use lance::dataset::{CommitBuilder, InsertBuilder, WriteParams, WriteProgressFn}; use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter; use lance_table::format::Fragment; @@ -74,7 +74,9 @@ fn merge_transactions(mut transactions: Vec) -> Option /// /// This plan executes inserts by: /// 1. Each partition writes data independently using InsertBuilder::execute_uncommitted_stream -/// 2. The last partition to complete commits all transactions atomically +/// 2. The last partition to complete merges transactions, optionally attaches one +/// precomputed generated-column metadata patch when the merged write has rows, +/// then commits once /// 3. Returns the count of inserted rows per partition #[derive(Debug)] pub struct InsertExec { @@ -83,6 +85,10 @@ pub struct InsertExec { input: Arc, write_params: WriteParams, tracker: Option>, + /// Optional whole-transaction field-metadata patch for generated-column + /// invalidation. Attached once after partition merge, and only when the + /// merged operation contains at least one written row. + schema_metadata_updates: Option, properties: Arc, partial_transactions: Arc>>, metrics: ExecutionPlanMetricsSet, @@ -95,7 +101,7 @@ impl InsertExec { input: Arc, write_params: WriteParams, ) -> Self { - Self::new_with_tracker(ds_wrapper, dataset, input, write_params, None) + Self::new_with_tracker(ds_wrapper, dataset, input, write_params, None, None) } pub(crate) fn new_with_tracker( @@ -104,6 +110,7 @@ impl InsertExec { input: Arc, write_params: WriteParams, tracker: Option>, + schema_metadata_updates: Option, ) -> Self { let schema = COUNT_SCHEMA.clone(); let num_partitions = input.output_partitioning().partition_count(); @@ -120,6 +127,7 @@ impl InsertExec { input, write_params, tracker, + schema_metadata_updates, properties: Arc::new(properties), partial_transactions: Arc::new(Mutex::new(Vec::with_capacity(num_partitions))), metrics: ExecutionPlanMetricsSet::new(), @@ -176,6 +184,7 @@ impl ExecutionPlan for InsertExec { children[0].clone(), self.write_params.clone(), self.tracker.clone(), + self.schema_metadata_updates.clone(), ))) } @@ -191,6 +200,7 @@ impl ExecutionPlan for InsertExec { let total_partitions = self.input.output_partitioning().partition_count(); let ds_wrapper = self.ds_wrapper.clone(); let tracker = self.tracker.clone(); + let schema_metadata_updates = self.schema_metadata_updates.clone(); let output_bytes = MetricBuilder::new(&self.metrics).output_bytes(partition); let input_schema = input_stream.schema(); @@ -220,6 +230,8 @@ impl ExecutionPlan for InsertExec { })); } + // Each partition stages an uncommitted data-only transaction. + // Metadata invalidation is attached once on the merged commit. let transaction = InsertBuilder::new(dataset.clone()) .with_params(&write_params) .execute_uncommitted_stream(input_stream) @@ -241,8 +253,15 @@ impl ExecutionPlan for InsertExec { }; if let Some(transactions) = to_commit - && let Some(merged_txn) = merge_transactions(transactions) + && let Some(mut merged_txn) = merge_transactions(transactions) { + // Attach the precomputed patch only for non-empty writes, and + // only once for the whole multi-partition transaction. + if count_rows_from_operation(&merged_txn.operation) > 0 + && let Some(updates) = schema_metadata_updates + { + merged_txn = merged_txn.with_schema_metadata_updates(updates)?; + } let new_dataset = CommitBuilder::new(dataset.clone()) .execute(merged_txn) .await?;