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