feat(dataset): per-fragment column writes that survive compaction

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<String> {
+    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<String> {
+    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<i32> {
+    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<ArrowField>| 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<Dataset>, 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/<file-stem>/ 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<Item = Result<RecordBatch>> + Send,
+        schema: &Schema,
+    ) -> Result<super::transaction::DataReplacementGroup> {
+        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::<Vec<_>>(),
+        );
+
+        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