From 0790f20d108bc94836e225043eccbebdcdcf65e1 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Fri, 7 Aug 2026 17:20:53 -0700 Subject: [PATCH] feat(dataset): per-fragment column writes that survive compaction Add FileFragment::write_column, which stages new data for one fragment's column as a standalone data file and returns its DataReplacementGroup without committing it. add_columns only appends a new field; write_column may stage a file that answers for a field the fragment already has, so a caller can recompute a column instead of only adding one. A field the dataset already defines has to match its manifest definition and be staged at the path the manifest gives it, its physical layout comes from the manifest rather than the caller, and each batch is checked against that tree before being reordered to match it. Extend DataReplacement to handle one more layout. Today it swaps a file when the field sets match exactly, or appends one when the fragment does not cover the fields at all, and rejects everything else. That rejection covers the layout a long-lived column actually reaches: compaction folds the column into a shared base file, no file's field set matches a single-column replacement any more, and nothing can replace the column again. Where the replaced fields all sit inside one wider file, tombstone them in place and append the new file to answer for them. Legacy V1 files keep exact-match replacement, because their reader derives page table offsets from the first field in the file metadata. Supersedes #8207, which proposed a Dataset-level staging and commit protocol for the same goal. Review there argued for building on the fragment-level API instead. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_017atn3ck33ob15HFVqxw54v --- rust/lance/src/dataset/fragment.rs | 301 +++++++- .../dataset/tests/fragment_write_column.rs | 683 ++++++++++++++++++ rust/lance/src/dataset/tests/mod.rs | 1 + rust/lance/src/dataset/transaction.rs | 177 ++++- 4 files changed, 1149 insertions(+), 13 deletions(-) create mode 100644 rust/lance/src/dataset/tests/fragment_write_column.rs diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 621754c55f0..94cfa6feab2 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -17,21 +17,23 @@ use arrow_array::types::UInt64Type; use arrow_array::{ Array, RecordBatch, RecordBatchReader, StructArray, UInt32Array, UInt64Array, new_null_array, }; -use arrow_schema::Schema as ArrowSchema; +use arrow_schema::{DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema}; use datafusion::logical_expr::Expr; use datafusion::scalar::ScalarValue; use futures::future::{BoxFuture, try_join_all}; -use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, join, stream}; +use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, join, stream}; use lance_arrow::json::{convert_json_columns, has_json_fields, is_arrow_json_field}; use lance_arrow::{RecordBatchExt, SchemaExt}; -use lance_core::datatypes::{OnMissing, OnTypeMismatch, SchemaCompareOptions}; +use lance_core::datatypes::{ + NullabilityComparison, OnMissing, OnTypeMismatch, SchemaCompareOptions, +}; use lance_core::utils::address::RowAddress; use lance_core::utils::deletion::DeletionVector; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::{ Error, Result, cache::{CacheKey, CacheKeySchema, KeyBuilder}, - datatypes::Schema, + datatypes::{Field, Schema, Schema as LanceSchema}, }; use lance_core::{ ROW_ADDR, ROW_ADDR_FIELD, ROW_CREATED_AT_VERSION_FIELD, ROW_ID, ROW_ID_FIELD, @@ -56,6 +58,7 @@ use lance_table::utils::stream::{ ReadBatchFutStream, ReadBatchTask, ReadBatchTaskStream, RowIdAndDeletesConfig, wrap_with_row_id_and_delete, }; +use object_store::path::Path; use roaring::RoaringBitmap; use self::write::FragmentCreateBuilder; @@ -65,7 +68,7 @@ use super::rowids::load_row_id_sequence; use super::scanner::Scanner; use super::updater::Updater; -use super::{NewColumnTransform, WriteParams, schema_evolution}; +use super::{NewColumnTransform, WriteParams, schema_evolution, versions}; use crate::dataset::Dataset; use crate::dataset::fragment::session::FragmentSession; use crate::dataset::overlay::{ @@ -732,6 +735,80 @@ enum MetadataMode { Full, } +/// The first path in `fields` that names a sibling twice. Projection picks +/// children by name, so a duplicate makes that choice arbitrary, and the +/// name-set comparison the schema check uses cannot see one at all. +fn duplicate_field_path(fields: &ArrowFields, path: &str) -> Option { + let mut seen = HashSet::new(); + for field in fields { + let qualified = if path.is_empty() { + field.name().clone() + } else { + format!("{path}.{}", field.name()) + }; + if !seen.insert(field.name()) { + return Some(qualified); + } + if let Some(nested) = duplicate_nested_path(field.data_type(), &qualified) { + return Some(nested); + } + } + None +} + +fn duplicate_nested_path(data_type: &DataType, path: &str) -> Option { + match data_type { + DataType::Struct(children) => duplicate_field_path(children, path), + DataType::List(item) + | DataType::LargeList(item) + | DataType::FixedSizeList(item, _) + | DataType::Map(item, _) => { + duplicate_nested_path(item.data_type(), &format!("{path}.item")) + } + _ => None, + } +} + +/// The first id in `field`'s subtree that `schema` already defines elsewhere. +fn borrowed_field_id(field: &Field, schema: &Schema) -> Option { + if schema.field_by_id(field.id).is_some() { + return Some(field.id); + } + field + .children + .iter() + .find_map(|child| borrowed_field_id(child, schema)) +} + +/// `field` with nullability dropped at every level: the projector rebuilds +/// arrays against its target and panics rather than reports on a constraint, +/// so it gets a shape that cannot fail and the writer objects instead. +fn relax_nullability(field: &ArrowField) -> ArrowField { + let relax = |field: &Arc| Arc::new(relax_nullability(field)); + let data_type = match field.data_type() { + DataType::Struct(children) => DataType::Struct(children.iter().map(relax).collect()), + DataType::List(item) => DataType::List(relax(item)), + DataType::LargeList(item) => DataType::LargeList(relax(item)), + DataType::FixedSizeList(item, width) => DataType::FixedSizeList(relax(item), *width), + // A Map's entries struct and its key stay required -- Arrow rejects a + // map whose entries or keys are nullable -- so only the value relaxes. + DataType::Map(entries, sorted) => match entries.data_type() { + DataType::Struct(kv) if kv.len() == 2 => { + let value = Arc::new(relax_nullability(&kv[1])); + let entries = ArrowField::new( + entries.name(), + DataType::Struct(vec![kv[0].clone(), value].into()), + false, + ); + DataType::Map(Arc::new(entries), *sorted) + } + _ => field.data_type().clone(), + }, + other => other.clone(), + }; + ArrowField::new(field.name(), data_type, true).with_metadata(field.metadata().clone()) +} + impl FileFragment { /// Creates a new FileFragment. pub fn new(dataset: Arc, metadata: Fragment) -> Self { @@ -2011,6 +2088,220 @@ impl FileFragment { Ok((fragments.into_iter().next().unwrap(), schema)) } + fn schema_mismatch(&self, detail: impl std::fmt::Display) -> Error { + Error::invalid_input(format!( + "column data for fragment {} does not match the requested schema: {detail}", + self.id() + )) + } + + /// Remove a staged file that will not be returned. Best effort: it is + /// unreachable either way, and must not mask the error that caused it. + async fn discard_staged_file(&self, path: &Path) { + // Blob v2 spills sidecars into data// beside the file, and + // those are the large ones; leaving them is what makes a routine + // rejection expensive. + if let Some(stem) = path + .filename() + .and_then(|name| name.strip_suffix(".lance")) + .map(|stem| self.dataset.data_dir().join(stem)) + && let Err(delete_error) = self.dataset.object_store.remove_dir_all(stem.clone()).await + { + log::warn!("failed to delete staged blob sidecars '{stem}': {delete_error}"); + } + if let Err(delete_error) = self.dataset.object_store.delete(path).await { + log::warn!("failed to delete staged column file '{path}': {delete_error}"); + } + } + + /// Write new data for a column of this fragment as a standalone data file, + /// without committing it, and return the + /// [`DataReplacementGroup`](super::transaction::DataReplacementGroup) + /// describing it. + /// + /// Unlike [`Self::add_columns`], the staged file may answer for a field that + /// already exists, so this recomputes a column rather than only appending. + /// + /// `schema` names the fields being written. A field the dataset already + /// defines must match its manifest definition and appear at the path the + /// manifest gives it; anything else is a new column, and nothing beneath it + /// may reuse an existing field's id. Physical layout comes from the + /// manifest, so staging cannot change a field's storage encoding. Batches + /// are matched by name at every level, so struct children may arrive in any + /// order, but a batch not describing exactly that tree is rejected. + /// + /// `data` must produce exactly the fragment's physical row count, nulls + /// included: the file is positionally aligned with the fragment and no + /// deletion vector is applied on the way in. Batches are pulled one at a + /// time, so the full column need not be held in memory. + /// + /// Staging does not bind the result to the version it was computed from. + /// Concurrent replacements of the same field conflict at commit, but only + /// for a caller that commits with the version it actually read; a caller + /// that supplies a newer one publishes stale values unchallenged. + pub async fn write_column( + &self, + data: impl Stream> + Send, + schema: &Schema, + ) -> Result { + let expected_rows = self.physical_rows().await? as u64; + + // Readers take everything but the field id from the manifest, so a + // staged field reusing an id is decoded as the manifest's version rather + // than rejected. Compare full identity, not just the storage type. + let compare_options = SchemaCompareOptions { + compare_field_ids: true, + ..Default::default() + }; + // Top-level requests match top-level manifest fields only: resolving an + // id from anywhere lets a caller reuse a field at a path the dataset + // never gave it, staging a file covering the borrowed field. Layout then + // comes from the manifest, since the metadata the identity check ignores + // -- packed structs, blob encoding -- decides physical field coverage. + let dataset_schema = self.dataset.schema(); + let mut writer_fields = Vec::with_capacity(schema.fields.len()); + for field in &schema.fields { + // The scanner injects these itself; a stored copy collides with it + // at projection time. Same boundary the ordinary insert path draws. + if lance_core::is_system_column(&field.name) { + return Err(Error::invalid_input(format!( + "column data for fragment {} names reserved column '{}'", + self.id(), + field.name + ))); + } + match dataset_schema + .fields + .iter() + .find(|existing| existing.id == field.id) + { + Some(existing) => { + // `explain_difference` recurses, covering the whole subtree. + if let Some(difference) = field.explain_difference(existing, &compare_options) { + return Err(Error::invalid_input(format!( + "column data for fragment {} does not match dataset field id {}: {}", + self.id(), + field.id, + difference + ))); + } + writer_fields.push(existing.clone()); + } + None => { + // A new column: nothing beneath it may reuse an existing id. + if let Some(borrowed) = borrowed_field_id(field, dataset_schema) { + return Err(Error::invalid_input(format!( + "column data for fragment {} puts dataset field id {} under new column '{}': \ + a field must be written at the path the dataset gives it", + self.id(), + borrowed, + field.name + ))); + } + writer_fields.push(field.clone()); + } + } + } + let writer_schema = Schema { + fields: writer_fields, + metadata: schema.metadata.clone(), + }; + let batch_schema = ArrowSchema::from(&writer_schema); + let projection_schema = ArrowSchema::new( + batch_schema + .fields() + .iter() + .map(|field| relax_nullability(field)) + .collect::>(), + ); + + let file_version = ConcreteFileVersion::from( + self.dataset + .manifest + .data_storage_format + .lance_file_version()?, + ); + + if file_version == ConcreteFileVersion::V1 { + // The legacy reader pairs a fragment's files by batch boundary, so a + // staged file chunked to the caller's batches leaves the fragment + // unreadable. Rechunking is the legacy update path's job, not this + // one's. + return Err(Error::not_supported(format!( + "write_column is not supported for fragment {} in the legacy file format", + self.id() + ))); + } + + // The update writer, not a raw file writer: that boundary carries the + // version's write policies (blob v2 columns arrive logical and must be + // prepared for the encoders) and returns a populated `DataFile`. + let mut writer = + versions::open_update_writer(file_version, self.dataset.as_ref(), &writer_schema) + .await?; + let staged_path = { + let (file_name, _) = writer.data_file_path(); + self.dataset.data_dir().join(file_name) + }; + + let mut data = std::pin::pin!(data); + while let Some(batch_result) = data.next().await { + let batch = batch_result?; + // Struct encoders consume children positionally, so a batch ordered + // differently from the manifest lands under the wrong field ids. + // Projection fixes that by name, but it downcasts by shape, so the + // whole tree is compared first. Nullability is the writer's to + // enforce, against the data rather than the declared schema. + if let Some(duplicate) = duplicate_field_path(batch.schema_ref().fields(), "") { + self.discard_staged_file(&staged_path).await; + return Err(self.schema_mismatch(format!("column '{duplicate}' appears twice"))); + } + if let Err(mismatch) = + LanceSchema::try_from(batch.schema_ref().as_ref()).and_then(|staged| { + staged.check_compatible( + &writer_schema, + &SchemaCompareOptions { + compare_nullability: NullabilityComparison::Ignore, + ignore_field_order: true, + ..Default::default() + }, + ) + }) + { + self.discard_staged_file(&staged_path).await; + return Err(self.schema_mismatch(mismatch)); + } + let batch = match batch.project_by_schema(&projection_schema) { + Ok(batch) => batch, + Err(err) => { + self.discard_staged_file(&staged_path).await; + return Err(self.schema_mismatch(err)); + } + }; + // The writer applies the manifest's nullability rule to the data; + // a batch it turns down leaves a file nothing will ever reference. + if let Err(err) = writer.write(std::slice::from_ref(&batch)).await { + self.discard_staged_file(&staged_path).await; + return Err(err); + } + } + let (num_rows, data_file) = writer.finish().await?; + if num_rows as u64 != expected_rows { + self.discard_staged_file(&staged_path).await; + return Err(Error::invalid_input(format!( + "column data for fragment {} has {} rows but the fragment has {} physical rows", + self.id(), + num_rows, + expected_rows + ))); + } + + Ok(super::transaction::DataReplacementGroup( + self.id() as u64, + data_file, + )) + } + /// Delete rows from the fragment. /// /// If all rows are deleted, returns `Ok(None)`. Otherwise, returns a new diff --git a/rust/lance/src/dataset/tests/fragment_write_column.rs b/rust/lance/src/dataset/tests/fragment_write_column.rs new file mode 100644 index 00000000000..35c794db92c --- /dev/null +++ b/rust/lance/src/dataset/tests/fragment_write_column.rs @@ -0,0 +1,683 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Per-fragment column writes: staging a column's data as a standalone file +//! with `FileFragment::write_column`, and committing it as a `DataReplacement` +//! whose coverage may not line up with any single file -- the case a computed +//! column reaches once compaction folds it into a shared base file. + +use std::sync::Arc; + +use arrow::array::AsArray; +use arrow_array::types::{Int32Type, UInt64Type}; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Int32Array, ListArray, RecordBatch, RecordBatchIterator, + StructArray, +}; +use arrow_buffer::OffsetBuffer; +use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; +use futures::{TryStreamExt, stream}; +use lance_core::datatypes::Schema as LanceSchema; +use lance_core::utils::tempfile::TempStrDir; +use lance_core::{ROW_ID, ROW_LAST_UPDATED_AT_VERSION}; +use lance_encoding::constants::PACKED_STRUCT_META_KEY; +use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; +use rstest::rstest; + +use crate::dataset::optimize::{CompactionOptions, compact_files}; +use crate::dataset::schema_evolution::NewColumnTransform; +use crate::dataset::transaction::{DataReplacementGroup, Operation}; +use crate::dataset::write::WriteParams; +use crate::dataset::{WriteDestination, fragment::FileFragment}; +use crate::{Dataset, Result}; + +fn batch_of(fields: Vec, columns: Vec) -> RecordBatch { + RecordBatch::try_new(Arc::new(ArrowSchema::new(fields)), columns).unwrap() +} + +fn ints(values: Vec) -> ArrayRef { + Arc::new(Int32Array::from(values)) as ArrayRef +} + +async fn dataset_of(batch: RecordBatch, version: Option) -> Dataset { + let schema = batch.schema(); + let params = version.map(|data_storage_version| WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }); + Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://", + params, + ) + .await + .unwrap() +} + +/// A one-fragment dataset holding a single non-null `id` column of `[1, 2]`. +async fn id_dataset() -> Dataset { + id_dataset_of(2, 1024).await +} + +fn only_fragment(dataset: &Dataset) -> FileFragment { + dataset.get_fragments().into_iter().next().unwrap() +} + +/// Lance schema for a new nullable Int32 column with a fresh field id. +fn new_column_schema(dataset: &Dataset, name: &str) -> LanceSchema { + let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new( + name, + DataType::Int32, + true, + )])) + .unwrap(); + schema.fields[0].id = dataset.manifest.max_field_id() + 1; + schema +} + +async fn stage( + dataset: &Dataset, + batch: RecordBatch, + schema: &LanceSchema, +) -> Result { + only_fragment(dataset) + .write_column(stream::iter([Ok(batch)]), schema) + .await +} + +async fn commit(dataset: &Dataset, replacements: Vec) -> Result { + let read_version = dataset.manifest.version; + Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset.clone())), + Operation::DataReplacement { replacements }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await +} + +/// A multi-fragment dataset of `rows` sequential ids, with stable row ids so +/// replacements can be checked against row lineage. +async fn id_dataset_of(rows: i32, max_rows_per_file: usize) -> Dataset { + let batch = batch_of( + vec![ArrowField::new("id", DataType::Int32, false)], + vec![ints((1..=rows).collect())], + ); + let schema = batch.schema(); + Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://", + Some(WriteParams { + max_rows_per_file, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap() +} + +async fn declare_all_null(dataset: &mut Dataset, name: &str) { + let arrow = Arc::new(ArrowSchema::new(vec![ArrowField::new( + name, + DataType::Int32, + true, + )])); + dataset + .add_columns(NewColumnTransform::AllNulls(arrow), None, None) + .await + .unwrap(); +} + +/// Stage `values` for an existing `column` of one fragment. +async fn stage_column( + dataset: &Dataset, + fragment_id: u64, + column: &str, + values: Vec, +) -> DataReplacementGroup { + let schema = LanceSchema { + fields: vec![dataset.schema().field(column).unwrap().clone()], + metadata: Default::default(), + }; + let batch = batch_of( + vec![ArrowField::new(column, DataType::Int32, true)], + vec![ints(values)], + ); + dataset + .get_fragments() + .into_iter() + .find(|fragment| fragment.id() as u64 == fragment_id) + .expect("fragment to stage for") + .write_column(stream::iter([Ok(batch)]), &schema) + .await + .unwrap() +} + +fn values(batch: &RecordBatch, name: &str) -> Vec> { + let col = batch[name].as_primitive::(); + (0..batch.num_rows()) + .map(|i| (!col.is_null(i)).then(|| col.value(i))) + .collect() +} + +/// A `point` struct of two non-null Int32 children, packed or not. +fn point_schema(packed: bool) -> Arc { + let mut point = ArrowField::new("point", DataType::Struct(point_children()), false); + if packed { + point.set_metadata([(PACKED_STRUCT_META_KEY.to_string(), "true".to_string())].into()); + } + Arc::new(ArrowSchema::new(vec![point])) +} + +fn point_children() -> Fields { + Fields::from(vec![ + ArrowField::new("x", DataType::Int32, false), + ArrowField::new("y", DataType::Int32, false), + ]) +} + +/// `xs` and `ys` are matched to `schema`'s children by name, so a schema that +/// orders them y-then-x still receives each child's own values. +fn points(schema: &Arc, xs: [i32; 2], ys: [i32; 2]) -> RecordBatch { + let DataType::Struct(children) = schema.field(0).data_type().clone() else { + unreachable!("point schema is a struct") + }; + let columns = children + .iter() + .map(|child| ints(if child.name() == "x" { xs } else { ys }.to_vec())) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StructArray::new(children, columns, None)) as ArrayRef], + ) + .unwrap() +} + +/// Commit `group` and read the `point` column back as its two children. +async fn committed_points(dataset: &Dataset, group: DataReplacementGroup) -> (Vec, Vec) { + let batch = commit(dataset, vec![group]) + .await + .unwrap() + .scan() + .try_into_batch() + .await + .unwrap(); + let child = |i: usize| { + batch + .column(0) + .as_struct() + .column(i) + .as_primitive::() + .values() + .to_vec() + }; + (child(0), child(1)) +} + +#[rstest] +#[tokio::test] +async fn test_records_writer_layout( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, +) { + let dataset = dataset_of( + arrow_array::record_batch!(("id", Int32, [1, 2])).unwrap(), + Some(version), + ) + .await; + let schema = new_column_schema(&dataset, "value"); + let fragment = only_fragment(&dataset); + + // Streamed as two batches: the DataFile must record the writer's + // field/column layout and the dataset's file version. + let DataReplacementGroup(replaced, data_file) = fragment + .write_column( + stream::iter([ + Ok(arrow_array::record_batch!(("value", Int32, [1])).unwrap()), + Ok(arrow_array::record_batch!(("value", Int32, [2])).unwrap()), + ]), + &schema, + ) + .await + .unwrap(); + + assert_eq!(replaced, fragment.id() as u64); + assert_eq!(data_file.fields.as_ref(), &[schema.fields[0].id]); + assert_eq!(data_file.fields.len(), data_file.column_indices.len()); + assert!(data_file.path.ends_with(".lance")); + assert_eq!( + (data_file.file_major_version, data_file.file_minor_version), + ConcreteFileVersion::from(version).to_data_file_numbers() + ); +} + +/// Input `write_column` turns down before anything can be committed. The +/// container cases matter twice over: projection reorders by name but downcasts +/// by shape, so an unchecked batch is dropped silently or panics. +#[rstest] +#[case::too_few_rows("short", "physical rows")] +#[case::too_many_rows("long", "physical rows")] +#[case::unrequested_column("extra", "unexpected=[unrequested]")] +#[case::wrong_container("struct", "should have type int32 but type was struct")] +#[case::reserved_system_name("rowid", "reserved column")] +// The reader takes type, nullability and nested layout from the manifest, so a +// staged field reusing an id but differing in any of them would be decoded as +// the manifest's version rather than rejected -- `validate()` would not notice. +#[case::field_type_mismatch("wrong_type", "does not match dataset field id")] +#[case::field_nullability_mismatch("wrong_nullability", "does not match dataset field id")] +// Projection picks children by name, so a duplicate makes the choice arbitrary. +// The schema check compares name sets and cannot see one. +#[case::duplicate_column("duplicate", "appears twice")] +#[tokio::test] +async fn test_rejects_bad_input(#[case] shape: &str, #[case] expected: &str) { + let dataset = id_dataset().await; + let value = ArrowField::new("value", DataType::Int32, true); + let mut schema = new_column_schema(&dataset, "value"); + + let values = match shape { + "short" => batch_of(vec![value], vec![ints(vec![7])]), + "long" => batch_of(vec![value], vec![ints(vec![7, 8, 9])]), + "extra" => batch_of( + vec![value, ArrowField::new("unrequested", DataType::Int32, true)], + vec![ints(vec![1, 2]), ints(vec![3, 4])], + ), + "struct" => { + let inner = Fields::from(vec![ArrowField::new("x", DataType::Int32, true)]); + batch_of( + vec![ArrowField::new( + "value", + DataType::Struct(inner.clone()), + true, + )], + vec![Arc::new(StructArray::new(inner, vec![ints(vec![1, 2])], None)) as ArrayRef], + ) + } + "rowid" => { + schema = new_column_schema(&dataset, ROW_ID); + batch_of( + vec![ArrowField::new(ROW_ID, DataType::Int32, true)], + vec![ints(vec![1, 2])], + ) + } + "duplicate" => batch_of( + vec![value.clone(), value], + vec![ints(vec![1, 2]), ints(vec![3, 4])], + ), + "wrong_type" | "wrong_nullability" => { + let existing = dataset.schema().field("id").unwrap(); + let staged = if shape == "wrong_type" { + ArrowField::new("id", DataType::Float32, existing.nullable) + } else { + ArrowField::new("id", DataType::Int32, !existing.nullable) + }; + schema = LanceSchema::try_from(&ArrowSchema::new(vec![staged])).unwrap(); + schema.fields[0].id = existing.id; + batch_of( + vec![ArrowField::new("id", DataType::Int32, true)], + vec![ints(vec![1, 2])], + ) + } + other => unreachable!("unknown case {other}"), + }; + + let err = stage(&dataset, values, &schema).await.unwrap_err(); + assert!( + err.to_string().contains(expected), + "expected '{expected}' in error, got: {err}" + ); +} + +/// Nested containers the projector would otherwise reshape or unwind on: a +/// fixed-size list of the wrong width silently becomes a different row count, +/// and nulls under a required item panic instead of erroring. +#[rstest] +#[case::fixed_size_list_reshape( + true, + "fixed_size_list:int32:2 but type was fixed_size_list:int32:4" +)] +#[case::nulls_under_required_item(false, "non-null")] +#[tokio::test] +async fn test_rejects_bad_nested_input(#[case] reshape: bool, #[case] expected: &str) { + let item = |nullable| Arc::new(ArrowField::new("item", DataType::Int32, nullable)); + let nest = |kind: DataType, values: ArrayRef| { + batch_of(vec![ArrowField::new("v", kind, true)], vec![values]) + }; + // Fixed-size list: the same eight values, four rows of two against two of + // four. List: four values, one of them null under a required item. + let fsl = |width: i32| { + let array = FixedSizeListArray::new(item(true), width, ints((1..=8).collect()), None); + nest( + DataType::FixedSizeList(item(true), width), + Arc::new(array) as ArrayRef, + ) + }; + let list = |values: Vec>, nullable| { + let array = ListArray::new( + item(nullable), + OffsetBuffer::new(vec![0, 2, 4].into()), + Arc::new(Int32Array::from(values)) as ArrayRef, + None, + ); + nest(DataType::List(item(nullable)), Arc::new(array) as ArrayRef) + }; + let (seed, staged) = if reshape { + (fsl(2), fsl(4)) + } else { + ( + list(vec![Some(1), Some(2), Some(3), Some(4)], false), + list(vec![Some(10), None, Some(30), Some(40)], true), + ) + }; + + let dataset = dataset_of(seed, None).await; + let schema = dataset.schema().clone(); + let err = stage(&dataset, staged, &schema).await.unwrap_err(); + assert!( + err.to_string().contains(expected), + "expected '{expected}' in error, got: {err}" + ); +} + +/// A field id the dataset defines has to be staged where the manifest puts it. +/// Wrapping a root field under a new parent reuses its id at a path the dataset +/// never gave it, and committing that would overwrite the root column. +#[tokio::test] +async fn test_rejects_field_id_at_wrong_path() { + let dataset = dataset_of( + arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(), + None, + ) + .await; + let root_id = dataset.schema().field("value").unwrap().id; + + let inner = Fields::from(vec![ArrowField::new("value", DataType::Int32, true)]); + let wrapper_field = ArrowField::new("wrapper", DataType::Struct(inner.clone()), true); + let mut wrapper = + LanceSchema::try_from(&ArrowSchema::new(vec![wrapper_field.clone()])).unwrap(); + wrapper.fields[0].id = dataset.manifest.max_field_id() + 1; + wrapper.fields[0].children[0].id = root_id; + + let values = batch_of( + vec![wrapper_field], + vec![Arc::new(StructArray::new(inner, vec![ints(vec![10, 20])], None)) as ArrayRef], + ); + let err = stage(&dataset, values, &wrapper).await.unwrap_err(); + assert!( + err.to_string().contains("at the path the dataset gives it"), + "expected wrong-path error, got: {err}" + ); +} + +/// Field metadata decides physical layout -- a packed struct is one column, an +/// unpacked one a column per child -- so a caller's metadata must not be able to +/// stage a file whose coverage describes a different field set. +#[tokio::test] +async fn test_takes_layout_from_manifest() { + let arrow_schema = point_schema(true); + let dataset = dataset_of( + points(&arrow_schema, [1, 2], [10, 20]), + Some(LanceFileVersion::V2_1), + ) + .await; + let packed_field_id = dataset.schema().field("point").unwrap().id; + + // Identical to the manifest field but for the packed marker, which the + // field-identity comparison does not look at. + let mut staged_schema = dataset.schema().clone(); + staged_schema.fields[0] + .metadata + .remove(PACKED_STRUCT_META_KEY); + assert!(!staged_schema.fields[0].is_packed_struct()); + + let group = stage( + &dataset, + points(&arrow_schema, [3, 4], [30, 40]), + &staged_schema, + ) + .await + .unwrap(); + // Unpacked, the file would cover x and y instead, and DataReplacement would + // see coverage the packed field never had. + assert_eq!(group.1.fields.as_ref(), &[packed_field_id]); + + assert_eq!( + committed_points(&dataset, group).await, + (vec![3, 4], vec![30, 40]) + ); +} + +/// Struct encoders consume children positionally, so a batch whose children are +/// ordered differently from the manifest would be written under the wrong field +/// ids. Batches are matched by name at every level instead. +#[tokio::test] +async fn test_reorders_struct_children_by_name() { + let dataset = dataset_of(points(&point_schema(false), [1, 2], [10, 20]), None).await; + + // Names its children y-then-x: written positionally, y's values land in x. + let reordered = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "point", + DataType::Struct(point_children().into_iter().rev().cloned().collect()), + false, + )])); + let schema = dataset.schema().clone(); + let group = stage(&dataset, points(&reordered, [30, 40], [300, 400]), &schema) + .await + .unwrap(); + + assert_eq!( + committed_points(&dataset, group).await, + (vec![30, 40], vec![300, 400]), + "each child keeps its own values" + ); +} + +/// Blob columns arrive logical and must be prepared into sidecars and +/// descriptors before the V2.2+ structural encoders accept them, which is what +/// the per-version update writer does. +#[tokio::test] +async fn test_stages_blob_column() { + use crate::blob::{BlobArrayBuilder, blob_field}; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![blob_field("blob", true)])); + let blobs = |values: [&[u8]; 2]| { + let mut builder = BlobArrayBuilder::new(2); + for value in values { + builder.push_bytes(value).unwrap(); + } + RecordBatch::try_new(arrow_schema.clone(), vec![builder.finish().unwrap()]).unwrap() + }; + let dataset = dataset_of(blobs([b"one", b"two"]), Some(LanceFileVersion::V2_2)).await; + + let schema = dataset.schema().clone(); + let group = stage(&dataset, blobs([b"three", b"four"]), &schema) + .await + .unwrap(); + assert!( + !group.1.fields.as_ref().is_empty(), + "staged file must cover the blob field" + ); +} + +/// The computed-column lifecycle: declare all null, backfill, compact, and +/// refresh again. The refresh after compaction is the case that previously +/// failed with "no changes were made". +#[tokio::test] +async fn test_replacement_survives_compaction() { + let mut dataset = id_dataset_of(4, 2).await; + declare_all_null(&mut dataset, "v").await; + let v_id = dataset.schema().field("v").unwrap().id; + + let frag_ids: Vec = dataset + .get_fragments() + .iter() + .map(|f| f.id() as u64) + .collect(); + let mut replacements = Vec::new(); + for (i, frag_id) in frag_ids.iter().enumerate() { + let base = i as i32 * 100; + replacements.push(stage_column(&dataset, *frag_id, "v", vec![base + 1, base + 2]).await); + } + let mut dataset = commit(&dataset, replacements).await.unwrap(); + + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + let files = dataset.get_fragments()[0].metadata().files.clone(); + assert_eq!(files.len(), 1, "compaction folded the column into one file"); + assert!(files[0].fields.len() > 1); + + // Refresh the compacted fragment, repeatedly: every round must land its own + // values and reuse the appended file rather than stacking another one. + let fragment_id = dataset.get_fragments()[0].id() as u64; + let rows = dataset.get_fragments()[0].physical_rows().await.unwrap(); + for round in 0..3i32 { + let refreshed: Vec = (0..rows as i32).map(|r| round * 1000 + r).collect(); + let replacement = stage_column(&dataset, fragment_id, "v", refreshed.clone()).await; + dataset = commit(&dataset, vec![replacement]).await.unwrap(); + dataset.validate().await.unwrap(); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!( + values(&batch, "v"), + refreshed.iter().map(|v| Some(*v)).collect::>() + ); + assert_eq!( + dataset.get_fragments()[0].metadata().files.len(), + 2, + "round {round} changed the file count" + ); + } + assert_eq!( + dataset.schema().field("v").unwrap().id, + v_id, + "field id preserved" + ); + + let files = dataset.get_fragments()[0].metadata().files.clone(); + let covering: Vec<&[i32]> = files + .iter() + .filter(|f| f.fields.contains(&v_id)) + .map(|f| f.fields.as_ref()) + .collect(); + assert_eq!(covering.as_slice(), &[[v_id].as_slice()]); + + // Tombstoning into a wider file has to advance row lineage like any other + // replacement, or a delta consumer never learns the refresh happened. + let version = dataset.version().version; + let batch = dataset + .scan() + .project(&["v", ROW_LAST_UPDATED_AT_VERSION]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!( + batch[ROW_LAST_UPDATED_AT_VERSION] + .as_primitive::() + .values(), + vec![version; rows].as_slice() + ); +} + +/// The existing uncovered (all-null backfill) and exact-match paths must be +/// unchanged. +#[tokio::test] +async fn test_existing_paths_unchanged() { + // Uncovered -> push. + let mut dataset = id_dataset_of(2, 1024).await; + declare_all_null(&mut dataset, "v").await; + let frag_id = dataset.get_fragments()[0].id() as u64; + let r = stage_column(&dataset, frag_id, "v", vec![10, 20]).await; + let dataset = commit(&dataset, vec![r]).await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!( + values(&dataset.scan().try_into_batch().await.unwrap(), "v"), + vec![Some(10), Some(20)] + ); + let files_after_first = dataset.get_fragments()[0].metadata().files.len(); + + // Exact match -> in-place swap, no new file. + let r = stage_column(&dataset, frag_id, "v", vec![30, 40]).await; + let dataset = commit(&dataset, vec![r]).await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!( + values(&dataset.scan().try_into_batch().await.unwrap(), "v"), + vec![Some(30), Some(40)] + ); + assert_eq!( + dataset.get_fragments()[0].metadata().files.len(), + files_after_first, + "exact-match replacement swaps in place rather than appending" + ); +} + +/// The legacy reader pairs a fragment's files by batch boundary, so a staged +/// file chunked to the caller's batches would leave the fragment unreadable. +#[tokio::test] +async fn test_rejects_legacy_format() { + let dataset = dataset_of( + arrow_array::record_batch!(("id", Int32, [1, 2])).unwrap(), + Some(LanceFileVersion::Legacy), + ) + .await; + let schema = new_column_schema(&dataset, "value"); + let batch = batch_of( + vec![ArrowField::new("value", DataType::Int32, true)], + vec![ints(vec![1, 2])], + ); + let err = stage(&dataset, batch, &schema).await.unwrap_err(); + assert!( + err.to_string().contains("legacy file format"), + "expected a legacy-format rejection, got: {err}" + ); +} + +/// Blob v2 spills sidecars into `data//`; a rejected stage that +/// leaves them behind orphans arbitrarily large objects. +#[tokio::test] +async fn test_discards_blob_sidecars_on_failure() { + use crate::blob::{BlobArrayBuilder, blob_field}; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![blob_field("blob", true)])); + let blobs = |count: usize| { + let mut builder = BlobArrayBuilder::new(count); + for _ in 0..count { + builder.push_bytes(vec![7u8; 128 * 1024]).unwrap(); + } + RecordBatch::try_new(arrow_schema.clone(), vec![builder.finish().unwrap()]).unwrap() + }; + let test_uri = TempStrDir::default(); + let dataset = Dataset::write( + RecordBatchIterator::new([Ok(blobs(2))], arrow_schema.clone()), + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let before = count_files(&dataset).await; + // Three rows against a two-row fragment: rejected only after the sidecars + // have been spilled. + let schema = dataset.schema().clone(); + stage(&dataset, blobs(3), &schema).await.unwrap_err(); + assert_eq!( + count_files(&dataset).await, + before, + "a rejected stage must not leave sidecars behind" + ); +} + +async fn count_files(dataset: &Dataset) -> usize { + dataset + .object_store + .read_dir_all(&dataset.data_dir(), None) + .try_fold(0usize, |count, _| async move { Ok(count + 1) }) + .await + .unwrap() +} diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs index 2c3aa203e15..18cf6c7fd1c 100644 --- a/rust/lance/src/dataset/tests/mod.rs +++ b/rust/lance/src/dataset/tests/mod.rs @@ -17,3 +17,4 @@ mod dataset_schema_evolution; mod dataset_transactions; mod dataset_versioning; mod fragment_validate_tombstones; +mod fragment_write_column; diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index c4d713e930a..d10173331dd 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -33,6 +33,7 @@ use lance_index::mem_wal::CompactedSsTable; use lance_index::{frag_reuse::FRAG_REUSE_INDEX_NAME, is_system_index}; use lance_io::object_store::ObjectStore; use lance_table::feature_flags::{FLAG_STABLE_ROW_IDS, apply_feature_flags}; +use lance_table::format::overlay::TOMBSTONE_FIELD_ID; use lance_table::rowids::read_row_ids; use lance_table::{ format::{ @@ -2370,6 +2371,7 @@ impl Transaction { // TODO(rmeng): check new file and fragment are the same length let mut columns_covered = HashSet::new(); + let mut replaced_in_place = false; for file in &mut new_frag.files { if file.fields == new_file.fields && file.file_major_version == new_file.file_major_version @@ -2379,6 +2381,7 @@ impl Transaction { file.path = new_file.path.clone(); file.file_size_bytes = new_file.file_size_bytes.clone(); file.base_id = new_file.base_id; + replaced_in_place = true; } columns_covered.extend(file.fields.iter()); } @@ -2391,6 +2394,51 @@ impl Transaction { .file_version() .expect("Expected valid file version"); new_frag.files.push(new_file.clone()); + } else if !replaced_in_place + && new_frag.files.iter().any(|file| { + !file.is_legacy_file() + && new_file + .fields + .iter() + .all(|field| file.fields.contains(field)) + }) + { + // The replaced fields all live inside one wider file -- + // compaction folds a column into a shared base file, say. + // Tombstone them there and append the new file to answer + // for them, the idiom `update_columns` uses. A partial + // overlap, or one spanning files, falls through to the + // error below: it describes no layout we can resolve. + // + // Legacy V1 is excluded: its reader derives the page table + // offset from the first field in the metadata, so + // tombstoning one field leaves its siblings decoding from + // the wrong pages. V1 keeps exact-match replacement. + new_file + .file_version() + .expect("Expected valid file version"); + for file in &mut new_frag.files { + // Same reason as the guard above. + if file.is_legacy_file() { + continue; + } + file.fields = file + .fields + .iter() + .map(|field| { + if new_file.fields.contains(field) { + TOMBSTONE_FIELD_ID + } else { + *field + } + }) + .collect::>() + .into(); + } + new_frag + .files + .retain(|file| file.fields.iter().any(|&f| f != TOMBSTONE_FIELD_ID)); + new_frag.files.push(new_file.clone()); } // Nothing changed in the current fragment, which is not expected -- error out @@ -2400,13 +2448,22 @@ impl Transaction { )); } - // New base values for these fields supersede any overlay - // still shadowing them; tombstone the overlaid fields so the - // replacement is not silently masked. + // New base values supersede any overlay still shadowing + // them, so tombstone the overlaid fields. An overlay + // committed after this transaction's snapshot is the newer + // value though -- the conflict resolver rebases these two + // precisely because the overlay wins -- so it stays, and + // being newer it stays last, preserving the ordering. + let (mut superseded, newer): (Vec<_>, Vec<_>) = new_frag + .overlays + .drain(..) + .partition(|overlay| overlay.committed_version <= self.read_version); lance_table::format::overlay::tombstone_overlay_fields( - &mut new_frag.overlays, + &mut superseded, &replaced_fields, ); + superseded.extend(newer); + new_frag.overlays = superseded; final_fragments.push(new_frag); } @@ -6892,9 +6949,10 @@ mod tests { #[test] fn test_data_replacement_tombstones_overlaid_fields() { // A DataReplacement writing new base values for field 5 must stop any - // overlay from shadowing those cells: field 5 is tombstoned in place + // overlay already shadowing those cells: field 5 is tombstoned in place // (preserving the overlay's field 3), and an overlay covering only field - // 5 is dropped entirely. + // 5 is dropped entirely. Both overlays predate the transaction's read + // version, which is what makes the replacement the newer value. let mut fragment = Fragment::new(0); fragment.files = vec![ DataFile::new_legacy_from_fields("f3.lance", vec![3], None), @@ -6907,12 +6965,12 @@ mod tests { roaring::RoaringBitmap::from_iter([0u32]), roaring::RoaringBitmap::from_iter([0u32]), ]), - committed_version: 3, + committed_version: 1, }, DataOverlayFile { data_file: DataFile::new_legacy_from_fields("o5.lance", vec![5], None), coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), - committed_version: 3, + committed_version: 1, }, ]; @@ -6953,6 +7011,109 @@ mod tests { assert_eq!(frag.overlays[0].data_file.fields.as_ref(), &[3, -2]); } + /// Replace field 5 in `fragment` at `read_version`, against a manifest at + /// `manifest_version`. + fn replace_field_5( + fragment: Fragment, + manifest_version: u64, + read_version: u64, + ) -> Result { + let schema = ArrowSchema::new(vec![ArrowField::new("v", DataType::Int32, true)]); + let mut manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + lance_table::format::DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + manifest.version = manifest_version; + + let txn = Transaction::new( + read_version, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new( + "v-new.lance", + vec![5], + vec![0], + ConcreteFileVersion::V2_0, + None, + None, + ), + )], + }, + None, + ); + txn.build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .map(|(manifest, _)| manifest.fragments[0].clone()) + } + + #[test] + fn test_data_replacement_rejects_subset_of_legacy_file() { + // The V1 reader derives its page table offset from the first field in + // the file metadata, so turning `[4, 5]` into `[-2, 5]` would leave + // field 4 decoding from field 5's pages. With no exact match to swap, + // the replacement must be rejected rather than corrupting the sibling. + let mut fragment = Fragment::new(0); + fragment.files = vec![DataFile::new_legacy_from_fields( + "wide.lance", + vec![4, 5], + None, + )]; + + let result = replace_field_5(fragment, 1, 1); + assert!( + result.is_err(), + "legacy subset replacement must be rejected, got: {:?}", + result.map(|fragment| fragment.files) + ); + } + + #[test] + fn test_data_replacement_preserves_overlay_newer_than_snapshot() { + // An overlay committed after this transaction read its snapshot holds + // the newer value; the conflict resolver rebases the two precisely + // because the overlay wins. Tombstoning it would discard a committed + // write, so only overlays the transaction could have seen are superseded. + let mut fragment = Fragment::new(0); + // One wider file, so the replacement takes the tombstone-and-append path. + fragment.files = vec![DataFile::new( + "wide.lance", + vec![4, 5], + vec![0, 1], + ConcreteFileVersion::V2_0, + None, + None, + )]; + fragment.overlays = vec![DataOverlayFile { + data_file: DataFile::new( + "newer.lance", + vec![5], + vec![0], + ConcreteFileVersion::V2_0, + None, + None, + ), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version: 7, + }]; + + // Staged against version 6, i.e. before the overlay landed. + let fragment = replace_field_5(fragment, 7, 6).unwrap(); + assert!(fragment.files.iter().any(|f| f.path == "v-new.lance")); + assert_eq!( + fragment.overlays.len(), + 1, + "overlay committed after the snapshot must survive" + ); + assert_eq!(fragment.overlays[0].data_file.fields.as_ref(), &[5]); + } + #[test] fn test_data_overlay_build_manifest_merges_duplicate_groups() { // Two groups targeting the same fragment must both survive (a HashMap