From 38df2a140e2c2624d0e25fa73ecf18a870544af9 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:56:11 +0000 Subject: [PATCH 1/5] fix: support blob v2 in fragment update_columns --- python/python/tests/test_fragment.py | 40 +++++++++++++++ rust/lance/src/dataset/fragment.rs | 54 ++++++++++++++++++-- rust/lance/src/dataset/schema_evolution.rs | 4 +- rust/lance/src/dataset/write/merge_insert.rs | 1 + 4 files changed, 92 insertions(+), 7 deletions(-) diff --git a/python/python/tests/test_fragment.py b/python/python/tests/test_fragment.py index 11276a2213d..468dfd76a7f 100644 --- a/python/python/tests/test_fragment.py +++ b/python/python/tests/test_fragment.py @@ -617,6 +617,46 @@ def test_fragment_update_columns_with_custom_join_key(tmp_path): assert result["name"][2] == "Chase" # id=3 should have name Chase +def test_fragment_update_columns_with_blob_v2(tmp_path): + data = pa.table( + { + "id": pa.array([1, 2, 3, 4]), + "payload": lance.blob_array([b"one", b"two", b"", None]), + } + ) + dataset_uri = tmp_path / "test_dataset_update_columns_blob_v2" + dataset = lance.write_dataset( + data, + dataset_uri, + data_storage_version="2.2", + ) + + fragment = dataset.get_fragment(0) + updated_fragment, fields_modified = fragment.update_columns( + pa.table( + { + "id": pa.array([2]), + "payload": lance.blob_array([b"NEW"]), + } + ), + left_on="id", + ) + + operation = LanceOperation.Update( + updated_fragments=[updated_fragment], + fields_modified=fields_modified, + ) + updated_dataset = LanceDataset.commit( + dataset_uri, + operation, + read_version=dataset.version, + ) + + result = updated_dataset.to_table(blob_handling="all_binary") + assert result["id"].to_pylist() == [1, 2, 3, 4] + assert result["payload"].to_pylist() == [b"one", b"NEW", b"", None] + + def test_fragment_update_columns_with_nulls(tmp_path): """Test fragment update columns with null values.""" # Create initial dataset diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 115d024c004..be947dff8c5 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -24,7 +24,7 @@ use futures::future::{BoxFuture, try_join_all}; use futures::{FutureExt, 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::{BlobHandling, 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; @@ -1737,11 +1737,15 @@ impl FileFragment { /// at a time. This can be useful to control memory usage when processing very large /// fields. The batch_size will only be used if the dataset is a v2 dataset. It will /// be ignored for v1 datasets. + /// + /// The `blob_handling` parameter controls the in-memory representation of blob + /// columns read by the updater. If unset, the dataset schema is used unchanged. pub(crate) async fn updater>( &self, columns: Option<&[T]>, schemas: Option<(Schema, Schema)>, batch_size: Option, + blob_handling: Option, ) -> Result { let mut schema = self.dataset.schema().clone(); @@ -1761,6 +1765,14 @@ impl FileFragment { schema = schema.project(&projection)?; } + if let Some(blob_handling) = blob_handling { + schema.fields = schema + .fields + .into_iter() + .map(|field| blob_handling.unload_if_needed(field)) + .collect(); + } + // If there is no projection, we at least need to read the row addresses with_row_addr |= !with_row_id && schema.fields.is_empty(); @@ -1830,7 +1842,7 @@ impl FileFragment { } pub(crate) async fn merge(mut self, join_column: &str, joiner: &HashJoiner) -> Result { - let mut updater = self.updater(Some(&[join_column]), None, None).await?; + let mut updater = self.updater(Some(&[join_column]), None, None, None).await?; while let Some(batch) = updater.next().await? { let batch = joiner @@ -1909,11 +1921,29 @@ impl FileFragment { if !read_columns.iter().any(|n| n.as_str() == ROW_ADDR) { read_columns.push(ROW_ADDR.to_string()); } + let descriptor_blob_ids = read_columns + .iter() + .filter_map(|column| self.schema().field(column)) + .filter(|field| field.is_blob_v2()) + .filter_map(|field| u32::try_from(field.id).ok()) + .collect::>(); + let has_blob_v2 = !descriptor_blob_ids.is_empty(); + let blob_handling = has_blob_v2.then(|| { + let materialized_blob_ids = self + .schema() + .fields_pre_order() + .filter(|field| field.is_blob()) + .filter_map(|field| u32::try_from(field.id).ok()) + .filter(|field_id| !descriptor_blob_ids.contains(field_id)) + .collect(); + BlobHandling::SomeBlobsBinary(materialized_blob_ids) + }); let mut updater = self .updater( Some(&read_columns), Some((write_schema.clone(), self.schema().clone())), None, + blob_handling, ) .await?; // Hash join: rows matched on the right-hand stream rewrite columns; track physical offsets via `_rowaddr`. @@ -1937,6 +1967,17 @@ impl FileFragment { )) })?; while let Some(batch) = updater.next().await? { + let batch = if has_blob_v2 { + crate::dataset::optimize::transform_blob_v2_batch( + &self.dataset, + self.schema(), + batch.clone(), + true, + ) + .await? + } else { + batch.clone() + }; let index_column = batch[left_on].clone(); let matched = joiner.matched_join_rows(index_column.clone())?; if let Some(addr_col) = batch.column_by_name(ROW_ADDR) { @@ -1952,7 +1993,7 @@ impl FileFragment { } } let updated_batch = joiner - .collect_with_fallback(batch, index_column, self.dataset()) + .collect_with_fallback(&batch, index_column, self.dataset()) .await?; updater.update(updated_batch).await?; } @@ -5548,7 +5589,10 @@ mod tests { let mut merged_fragments = Vec::new(); for fragment_id in fragment_ids { let fragment = &mut dataset.get_fragment(fragment_id).unwrap(); - let mut updater = fragment.updater(Some(&["i"]), None, None).await.unwrap(); + let mut updater = fragment + .updater(Some(&["i"]), None, None, None) + .await + .unwrap(); while let Some(batch) = updater.next().await.unwrap() { let input_col = batch.column_by_name("i").unwrap(); let result_col = mul(input_col, &Int32Array::new_scalar(2)).unwrap(); @@ -5803,7 +5847,7 @@ mod tests { let fragment = dataset.get_fragments().pop().unwrap(); // Write batch_s using add_columns - let mut updater = fragment.updater(Some(&["i"]), None, None).await?; + let mut updater = fragment.updater(Some(&["i"]), None, None, None).await?; updater.next().await?; updater.update(batch_s.clone()).await?; let frag = updater.finish().await?; diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 54943594d9e..3749d950ea9 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -536,7 +536,7 @@ async fn add_columns_impl( } let mut updater = match fragment - .updater(read_columns_ref, schemas.clone(), batch_size) + .updater(read_columns_ref, schemas.clone(), batch_size, None) .await { Ok(updater) => updater, @@ -614,7 +614,7 @@ async fn add_columns_from_stream( let mut last_seen_batch: Option = None; for fragment in fragments { let mut updater = match fragment - .updater::(Some(&[]), schemas.clone(), batch_size) + .updater::(Some(&[]), schemas.clone(), batch_size, None) .await { Ok(updater) => updater, diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 0c4cd120568..181c3237fcd 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -1571,6 +1571,7 @@ impl MergeInsertJob { Some(&read_columns), Some((write_schema, dataset.schema().clone())), None, + None, ) .await?; From 1014ee202a2754844a6b3bb74b4e77e8fe8db1ef Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 11:45:51 +0000 Subject: [PATCH 2/5] fix: handle nested and external blob updates --- python/python/tests/test_fragment.py | 93 +++++++ rust/lance/src/dataset/blob.rs | 140 +++++++++- rust/lance/src/dataset/fragment.rs | 25 +- rust/lance/src/dataset/optimize.rs | 372 ++++++++++++++++++------- rust/lance/src/dataset/updater.rs | 17 +- rust/lance/src/dataset/versions/mod.rs | 7 +- rust/lance/src/dataset/write.rs | 2 + 7 files changed, 555 insertions(+), 101 deletions(-) diff --git a/python/python/tests/test_fragment.py b/python/python/tests/test_fragment.py index 468dfd76a7f..deb48c57076 100644 --- a/python/python/tests/test_fragment.py +++ b/python/python/tests/test_fragment.py @@ -657,6 +657,99 @@ def test_fragment_update_columns_with_blob_v2(tmp_path): assert result["payload"].to_pylist() == [b"one", b"NEW", b"", None] +def test_fragment_update_columns_with_nested_blob_v2(tmp_path): + def info_array(names, payloads): + fields = [pa.field("name", pa.string()), lance.blob_field("blob")] + return pa.StructArray.from_arrays( + [pa.array(names), lance.blob_array(payloads)], fields=fields + ) + + dataset_uri = tmp_path / "test_dataset_update_columns_nested_blob_v2" + dataset = lance.write_dataset( + pa.table( + { + "id": pa.array([1, 2]), + "info": info_array(["a", "b"], [b"one", b"two"]), + } + ), + dataset_uri, + data_storage_version="2.2", + ) + + updated_fragment, fields_modified = dataset.get_fragment(0).update_columns( + pa.table( + { + "id": pa.array([2]), + "info": info_array(["B"], [b"NEW"]), + } + ), + left_on="id", + ) + updated_dataset = LanceDataset.commit( + dataset_uri, + LanceOperation.Update( + updated_fragments=[updated_fragment], + fields_modified=fields_modified, + ), + read_version=dataset.version, + ) + + info = updated_dataset.to_table(blob_handling="all_binary")["info"].combine_chunks() + assert info.field("name").to_pylist() == ["a", "B"] + assert info.field("blob").to_pylist() == [b"one", b"NEW"] + + +def test_fragment_update_columns_preserves_external_blob_v2(tmp_path): + dataset_uri = tmp_path / "test_dataset_update_columns_external_blob_v2" + external = tmp_path / "existing-payload.bin" + external.write_bytes(b"outside") + dataset = lance.write_dataset( + pa.table( + { + "id": pa.array([1, 2]), + "payload": lance.blob_array([external.as_uri(), b"two"]), + } + ), + dataset_uri, + data_storage_version="2.2", + allow_external_blob_outside_bases=True, + ) + + updated_fragment, fields_modified = dataset.get_fragment(0).update_columns( + pa.table( + { + "id": pa.array([2]), + "payload": lance.blob_array([b"NEW"]), + } + ), + left_on="id", + ) + updated_dataset = LanceDataset.commit( + dataset_uri, + LanceOperation.Update( + updated_fragments=[updated_fragment], + fields_modified=fields_modified, + ), + read_version=dataset.version, + ) + + result = updated_dataset.to_table(blob_handling="all_binary") + assert result["payload"].to_pylist() == [b"outside", b"NEW"] + + new_external = tmp_path / "new-payload.bin" + new_external.write_bytes(b"new outside") + with pytest.raises(ValueError, match="outside registered external bases"): + updated_dataset.get_fragment(0).update_columns( + pa.table( + { + "id": pa.array([2]), + "payload": lance.blob_array([new_external.as_uri()]), + } + ), + left_on="id", + ) + + def test_fragment_update_columns_with_nulls(tmp_path): """Test fragment update columns with null values.""" # Create initial dataset diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index b59250d670f..d2e94b478d8 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::{ - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, HashSet}, future::Future, ops::{DerefMut, Range}, panic::AssertUnwindSafe, @@ -14,7 +14,7 @@ use arrow::datatypes::{UInt8Type, UInt32Type, UInt64Type}; use arrow_array::{ Array, ArrayRef, GenericListArray, OffsetSizeTrait, RecordBatch, builder::LargeBinaryBuilder, }; -use arrow_buffer::{OffsetBuffer, ScalarBuffer}; +use arrow_buffer::{ArrowNativeType, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use bytes::Bytes; use futures::future::BoxFuture; @@ -188,6 +188,140 @@ impl ExternalBaseResolver { } } +fn collect_external_blob_uris( + field: &ArrowField, + array: &ArrayRef, + selected_rows: &[bool], + field_path: &str, + external_uris: &mut Vec<(String, String)>, +) -> Result<()> { + if array.len() != selected_rows.len() { + return Err(Error::internal(format!( + "Blob field '{}' row count {} did not match selection length {}", + field_path, + array.len(), + selected_rows.len() + ))); + } + + if field.is_blob_v2() { + let struct_array = array.as_struct(); + if BlobV2Layout::classify(struct_array.fields()) != Some(BlobV2Layout::Logical) { + return Err(blob_v2_shape_error(field, &[BlobV2Layout::Logical])); + } + let uri_column = struct_array + .column_by_name("uri") + .ok_or_else(|| Error::invalid_input("Blob struct missing `uri` field"))? + .as_string::(); + for (row_idx, is_selected) in selected_rows.iter().copied().enumerate() { + if is_selected && struct_array.is_valid(row_idx) && uri_column.is_valid(row_idx) { + external_uris.push(( + field_path.to_string(), + uri_column.value(row_idx).to_string(), + )); + } + } + return Ok(()); + } + + match field.data_type() { + ArrowDataType::Struct(children) => { + let struct_array = array.as_struct(); + let child_selection = selected_rows + .iter() + .copied() + .enumerate() + .map(|(row_idx, is_selected)| is_selected && struct_array.is_valid(row_idx)) + .collect::>(); + for (child_field, child_array) in children.iter().zip(struct_array.columns()) { + let child_path = format!("{}.{}", field_path, child_field.name()); + collect_external_blob_uris( + child_field, + child_array, + &child_selection, + &child_path, + external_uris, + )?; + } + } + ArrowDataType::List(child) => { + let list_array = array.as_list::(); + let mut child_selection = vec![false; list_array.values().len()]; + for (row_idx, is_selected) in selected_rows.iter().copied().enumerate() { + if is_selected && list_array.is_valid(row_idx) { + let start = list_array.value_offsets()[row_idx].as_usize(); + let end = list_array.value_offsets()[row_idx + 1].as_usize(); + child_selection[start..end].fill(true); + } + } + let child_path = format!("{}.{}", field_path, child.name()); + collect_external_blob_uris( + child, + list_array.values(), + &child_selection, + &child_path, + external_uris, + )?; + } + ArrowDataType::LargeList(child) => { + let list_array = array.as_list::(); + let mut child_selection = vec![false; list_array.values().len()]; + for (row_idx, is_selected) in selected_rows.iter().copied().enumerate() { + if is_selected && list_array.is_valid(row_idx) { + let start = list_array.value_offsets()[row_idx].as_usize(); + let end = list_array.value_offsets()[row_idx + 1].as_usize(); + child_selection[start..end].fill(true); + } + } + let child_path = format!("{}.{}", field_path, child.name()); + collect_external_blob_uris( + child, + list_array.values(), + &child_selection, + &child_path, + external_uris, + )?; + } + _ => {} + } + Ok(()) +} + +/// Validate external blob references supplied by selected input rows. +/// +/// Existing rows can contain trusted absolute references that were accepted by an earlier write. +/// Update paths use this check before allowing those fallback values through the writer, so newly +/// matched values must still resolve beneath a registered external base. +pub(super) async fn validate_external_blob_references( + resolver: &ExternalBaseResolver, + batch: &RecordBatch, + selected_rows: &[bool], +) -> Result<()> { + let mut external_uris = Vec::new(); + for (field, array) in batch.schema().fields().iter().zip(batch.columns()) { + collect_external_blob_uris( + field, + array, + selected_rows, + field.name(), + &mut external_uris, + )?; + } + + let mut validated_uris = HashSet::new(); + for (field_path, uri) in external_uris { + if validated_uris.insert(uri.clone()) + && resolver.resolve_external_uri(&uri).await?.is_none() + { + return Err(Error::invalid_input(format!( + "External blob URI '{}' in field '{}' is outside registered external bases (dataset root is not allowed)", + uri, field_path + ))); + } + } + Ok(()) +} + struct RollingPackedBlobWriter { current: Option, current_size: usize, @@ -3304,7 +3438,7 @@ async fn materialize_blob_v2_list_array( Ok(Arc::new(list_array)) } -async fn materialize_blob_v2_descriptors( +pub(super) async fn materialize_blob_v2_descriptors( dataset: &Arc, blob_field_id: u32, descriptions: &StructArray, diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index be947dff8c5..b1a419a0129 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -1921,9 +1921,15 @@ impl FileFragment { if !read_columns.iter().any(|n| n.as_str() == ROW_ADDR) { read_columns.push(ROW_ADDR.to_string()); } - let descriptor_blob_ids = read_columns + let selected_field_ids = read_columns .iter() .filter_map(|column| self.schema().field(column)) + .map(|field| field.id) + .collect::>(); + let descriptor_blob_ids = self + .schema() + .project_by_ids(&selected_field_ids, true) + .fields_pre_order() .filter(|field| field.is_blob_v2()) .filter_map(|field| u32::try_from(field.id).ok()) .collect::>(); @@ -1946,6 +1952,19 @@ impl FileFragment { blob_handling, ) .await?; + if has_blob_v2 { + updater.allow_external_blob_outside_bases(); + } + let external_base_resolver = if has_blob_v2 { + super::write::blob_v2_external_base_resolver( + Some(self.dataset()), + &WriteParams::default(), + &write_schema, + ) + .await? + } else { + None + }; // Hash join: rows matched on the right-hand stream rewrite columns; track physical offsets via `_rowaddr`. // Convert Arrow JSON columns (Utf8) to Lance JSON (LargeBinary) in the right stream // so they match the physical storage format read from the fragment's left batch. @@ -1995,6 +2014,10 @@ impl FileFragment { let updated_batch = joiner .collect_with_fallback(&batch, index_column, self.dataset()) .await?; + if let Some(resolver) = external_base_resolver.as_deref() { + super::blob::validate_external_blob_references(resolver, &updated_batch, &matched) + .await?; + } updater.update(updated_batch).await?; } diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 120a61c2fe8..7476e3376a9 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -102,17 +102,21 @@ use crate::index::DatasetIndexExt; use crate::io::commit::{DEFAULT_COMMIT_RETRY_TIMEOUT, commit_transaction, migrate_fragments}; use arrow::array::AsArray; use arrow::datatypes::{UInt8Type, UInt32Type, UInt64Type}; -use arrow_array::Array; -use arrow_array::RecordBatch; -use arrow_array::StructArray; use arrow_array::builder::{LargeBinaryBuilder, PrimitiveBuilder, StringBuilder}; -use arrow_buffer::NullBuffer; +use arrow_array::{ + Array, ArrayRef, GenericListArray, OffsetSizeTrait, RecordBatch, StructArray, UInt32Array, +}; +use arrow_buffer::{OffsetBuffer, ScalarBuffer}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField}; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::{StreamExt, TryStreamExt}; +use futures::future::BoxFuture; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use lance_arrow::{list::ListArrayExt, r#struct::StructArrayExt}; use lance_core::Error; use lance_core::datatypes::{ BLOB_V2_LOGICAL_FIELDS, BLOB_V2_LOGICAL_TYPE, BlobHandling, BlobKind, BlobV2Layout, + Field as LanceField, }; use lance_core::utils::tokio::get_num_compute_intensive_cpus; use lance_core::utils::tracing::{DATASET_COMPACTING_EVENT, TRACE_DATASET_EVENTS}; @@ -1054,22 +1058,21 @@ impl<'a> BlobV2Descriptor<'a> { } } -/// Result of row classification for blob v2 compaction. +/// Result of row classification for a blob v2 rewrite. struct RowClassification { row_classes: Vec, - blob_read_addrs: Vec, + data_blob_indices: Vec, } /// Classify each row of a blob v2 column as Null, External, or DataBlob. fn classify_rows( struct_arr: &StructArray, descriptor: &BlobV2Descriptor<'_>, - row_addrs: &arrow::array::UInt64Array, column_name: &str, ) -> Result { let num_rows = struct_arr.len(); let mut row_classes = Vec::with_capacity(num_rows); - let mut blob_read_addrs = Vec::with_capacity(num_rows); + let mut data_blob_indices = Vec::with_capacity(num_rows); for i in 0..num_rows { if struct_arr.is_null(i) || descriptor.kind_col.is_null(i) { @@ -1085,14 +1088,14 @@ fn classify_rows( row_classes.push(RowClass::External); } else { row_classes.push(RowClass::DataBlob); - blob_read_addrs.push(row_addrs.value(i)); + data_blob_indices.push(i); } } } Ok(RowClassification { row_classes, - blob_read_addrs, + data_blob_indices, }) } @@ -1102,19 +1105,57 @@ fn classify_rows( /// payloads in memory at once. async fn descriptor_to_logical_blob_array( dataset: &Arc, + blob_field_id: u32, + struct_arr: &StructArray, descriptor: &BlobV2Descriptor<'_>, classification: &RowClassification, - column_name: &str, - num_rows: usize, - null_buffer: Option, + row_addrs: &[u64], + field_name: &str, ) -> Result { - let blob_files = if classification.blob_read_addrs.is_empty() { - Vec::new() + if struct_arr.len() != row_addrs.len() { + return Err(Error::internal(format!( + "Blob v2 field '{}' row count {} did not match row address count {}", + field_name, + struct_arr.len(), + row_addrs.len() + ))); + } + let blob_data = if classification.data_blob_indices.is_empty() { + None } else { - super::blob::take_blobs_by_addresses(dataset, &classification.blob_read_addrs, column_name) - .await? + let indices = classification + .data_blob_indices + .iter() + .map(|index| { + u32::try_from(*index).map_err(|_| { + Error::internal(format!( + "Blob v2 row index {} in field '{}' does not fit in u32", + index, field_name + )) + }) + }) + .collect::>>()?; + let indices = UInt32Array::from(indices); + let descriptions = arrow_select::take::take(struct_arr, &indices, None)?; + let descriptions = descriptions.as_struct(); + let data_row_addrs = classification + .data_blob_indices + .iter() + .map(|index| row_addrs[*index]) + .collect::>(); + Some( + super::blob::materialize_blob_v2_descriptors( + dataset, + blob_field_id, + descriptions, + &data_row_addrs, + ) + .await?, + ) }; + let num_rows = struct_arr.len(); + let data_values = blob_data.as_ref().map(|data| data.as_binary::()); let mut data_builder = LargeBinaryBuilder::with_capacity(num_rows, 0); let mut uri_builder = StringBuilder::with_capacity(num_rows, 0); let mut out_position_builder = PrimitiveBuilder::::with_capacity(num_rows); @@ -1140,7 +1181,7 @@ async fn descriptor_to_logical_blob_array( let base = dataset.manifest().base_paths.get(&base_id).ok_or_else(|| { Error::internal(format!( "External blob in column '{}' references unknown base_id {}", - column_name, base_id + field_name, base_id )) })?; let absolute_uri = format!("{}/{}", base.path.trim_end_matches('/'), uri_val); @@ -1158,15 +1199,21 @@ async fn descriptor_to_logical_blob_array( } } RowClass::DataBlob => { - let blob_file = blob_files[blob_file_idx].as_ref().ok_or_else(|| { + let data_values = data_values.as_ref().ok_or_else(|| { Error::internal(format!( - "Non-null blob row {} in column '{}' resolved to null", - i, column_name + "Non-null blob row {} in field '{}' produced no materialized data", + i, field_name )) })?; - let data = blob_file.read().await?; + if data_values.is_null(blob_file_idx) { + return Err(Error::internal(format!( + "Non-null blob row {} in field '{}' resolved to null", + i, field_name + ))); + } + let data = data_values.value(blob_file_idx); blob_file_idx += 1; - data_builder.append_value(data.as_ref()); + data_builder.append_value(data); uri_builder.append_null(); out_position_builder.append_null(); out_size_builder.append_null(); @@ -1182,10 +1229,199 @@ async fn descriptor_to_logical_blob_array( Arc::new(out_position_builder.finish()), Arc::new(out_size_builder.finish()), ], - null_buffer, + struct_arr.nulls().cloned(), )?) } +fn transformed_arrow_field(field: &LanceField, data_type: ArrowDataType) -> Arc { + let arrow_field = ArrowField::from(field); + Arc::new( + ArrowField::new(arrow_field.name(), data_type, arrow_field.is_nullable()) + .with_metadata(arrow_field.metadata().clone()), + ) +} + +fn transform_blob_v2_array<'a>( + dataset: &'a Arc, + field: &'a LanceField, + array: ArrayRef, + row_addrs: Arc<[u64]>, +) -> BoxFuture<'a, Result<(ArrayRef, Arc)>> { + async move { + if field.is_blob_v2() { + let struct_arr = array.as_struct(); + return match BlobV2Layout::classify(struct_arr.fields()) { + Some(BlobV2Layout::Logical) => { + let output_field = transformed_arrow_field(field, array.data_type().clone()); + Ok((array, output_field)) + } + Some(BlobV2Layout::Descriptor) => { + let descriptor = BlobV2Descriptor::try_from_struct(struct_arr, &field.name)?; + let classification = classify_rows(struct_arr, &descriptor, &field.name)?; + let logical = descriptor_to_logical_blob_array( + dataset, + u32::try_from(field.id).map_err(|_| { + Error::internal(format!( + "Blob v2 field id {} for '{}' does not fit in u32", + field.id, field.name + )) + })?, + struct_arr, + &descriptor, + &classification, + row_addrs.as_ref(), + &field.name, + ) + .await?; + Ok(( + Arc::new(logical) as ArrayRef, + transformed_arrow_field(field, BLOB_V2_LOGICAL_TYPE.clone()), + )) + } + Some(actual) => Err(Error::invalid_input(format!( + "Blob v2 field '{}' has {actual} layout; expected logical or descriptor layout during rewrite", + field.name + ))), + None => Err(Error::invalid_input(format!( + "Blob v2 field '{}' has unrecognized layout {:?}; expected logical or descriptor layout during rewrite", + field.name, + struct_arr.fields() + ))), + }; + } + + match field.data_type() { + ArrowDataType::Struct(_) => { + let struct_arr = array.as_struct().normalize_slicing()?; + let parent_nulls = struct_arr.nulls().cloned(); + let struct_arr = struct_arr.pushdown_nulls()?; + if field.children.len() != struct_arr.num_columns() { + return Err(Error::internal(format!( + "Struct field '{}' expected {} children during blob rewrite, got {}", + field.name, + field.children.len(), + struct_arr.num_columns() + ))); + } + let mut child_arrays = Vec::with_capacity(field.children.len()); + let mut child_fields = Vec::with_capacity(field.children.len()); + for (child_field, child_array) in + field.children.iter().zip(struct_arr.columns()) + { + let (child_array, child_field) = transform_blob_v2_array( + dataset, + child_field, + child_array.clone(), + row_addrs.clone(), + ) + .await?; + child_arrays.push(child_array); + child_fields.push(child_field); + } + let struct_array = StructArray::try_new( + child_fields.clone().into(), + child_arrays, + parent_nulls, + )?; + let output_field = transformed_arrow_field( + field, + ArrowDataType::Struct(child_fields.into()), + ); + Ok((Arc::new(struct_array), output_field)) + } + ArrowDataType::List(_) => { + transform_blob_v2_list_array::( + dataset, + field, + array.as_list::(), + row_addrs, + ) + .await + } + ArrowDataType::LargeList(_) => { + transform_blob_v2_list_array::( + dataset, + field, + array.as_list::(), + row_addrs, + ) + .await + } + _ => { + let output_field = transformed_arrow_field(field, array.data_type().clone()); + Ok((array, output_field)) + } + } + } + .boxed() +} + +async fn transform_blob_v2_list_array( + dataset: &Arc, + field: &LanceField, + list_array: &GenericListArray, + row_addrs: Arc<[u64]>, +) -> Result<(ArrayRef, Arc)> { + let list_array = if list_array.null_count() > 0 { + list_array.filter_garbage_nulls() + } else { + list_array.clone() + }; + let offsets = list_array.value_offsets(); + let values_start = offsets[0].as_usize(); + let values_end = offsets[list_array.len()].as_usize(); + if values_end < values_start { + return Err(Error::internal(format!( + "List field '{}' has invalid offsets during blob rewrite", + field.name + ))); + } + + let values_len = values_end - values_start; + let mut normalized_offsets = Vec::with_capacity(list_array.len() + 1); + normalized_offsets.push(O::usize_as(0)); + let mut child_row_addrs = Vec::with_capacity(values_len); + for row_idx in 0..list_array.len() { + let start = offsets[row_idx].as_usize(); + let end = offsets[row_idx + 1].as_usize(); + if end < start { + return Err(Error::internal(format!( + "List field '{}' has decreasing offsets during blob rewrite", + field.name + ))); + } + let row_addr = row_addrs.get(row_idx).copied().ok_or_else(|| { + Error::internal(format!( + "List field '{}' row address count {} did not match row count {}", + field.name, + row_addrs.len(), + list_array.len() + )) + })?; + child_row_addrs.extend(std::iter::repeat_n(row_addr, end - start)); + normalized_offsets.push(O::usize_as(end - values_start)); + } + + let child = field.children.first().ok_or_else(|| { + Error::internal(format!( + "List field '{}' missing child during blob rewrite", + field.name + )) + })?; + let values = list_array.values().slice(values_start, values_len); + let (values, child_field) = + transform_blob_v2_array(dataset, child, values, Arc::<[u64]>::from(child_row_addrs)) + .await?; + let list_array = GenericListArray::::try_new( + child_field, + OffsetBuffer::new(ScalarBuffer::from(normalized_offsets)), + values, + list_array.nulls().cloned(), + )?; + let output_field = transformed_arrow_field(field, list_array.data_type().clone()); + Ok((Arc::new(list_array), output_field)) +} + pub(crate) async fn transform_blob_v2_batch( dataset: &Arc, schema: &lance_core::datatypes::Schema, @@ -1207,7 +1443,14 @@ pub(crate) async fn transform_blob_v2_batch( )) })? .0; - let row_addrs = batch.column(row_addr_idx).as_primitive::(); + let row_addrs: Arc<[u64]> = batch + .column(row_addr_idx) + .as_primitive::() + .values() + .iter() + .copied() + .collect::>() + .into(); let mut new_columns: Vec> = Vec::new(); let mut new_fields: Vec> = Vec::new(); @@ -1218,81 +1461,20 @@ pub(crate) async fn transform_blob_v2_batch( continue; } - let lance_field = schema.field(field.name()); - let is_blob_v2 = lance_field.is_some_and(|f| f.is_blob_v2()); - - if !is_blob_v2 { + let Some(lance_field) = schema.field(field.name()) else { new_columns.push(batch.column(col_idx).clone()); new_fields.push(field.clone()); continue; - } - - let struct_arr = batch - .column(col_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::internal(format!( - "Blob v2 column '{}' expected StructArray, got {:?}", - field.name(), - batch.column(col_idx).data_type() - )) - })?; - - match BlobV2Layout::classify(struct_arr.fields()) { - // Merge-insert may supply a logical blob v2 value directly from the - // source. It does not refer to a row in the target dataset. - Some(BlobV2Layout::Logical) => { - new_columns.push(batch.column(col_idx).clone()); - new_fields.push(field.clone()); - continue; - } - Some(BlobV2Layout::Descriptor) => {} - Some(actual) => { - return Err(Error::invalid_input(format!( - "Blob v2 column '{}' has {actual} layout; expected logical or descriptor layout during compaction", - field.name() - ))); - } - None => { - return Err(Error::invalid_input(format!( - "Blob v2 column '{}' has unrecognized layout {:?}; expected logical or descriptor layout during compaction", - field.name(), - struct_arr.fields() - ))); - } - } - - let column_name = field.name(); - let descriptor = BlobV2Descriptor::try_from_struct(struct_arr, column_name)?; - let classification = classify_rows(struct_arr, &descriptor, row_addrs, column_name)?; - let num_rows = struct_arr.len(); - - let new_struct = descriptor_to_logical_blob_array( + }; + let (new_column, new_field) = transform_blob_v2_array( dataset, - &descriptor, - &classification, - column_name, - num_rows, - struct_arr.nulls().cloned(), + lance_field, + batch.column(col_idx).clone(), + row_addrs.clone(), ) .await?; - - new_columns.push(Arc::new(new_struct)); - let logical_field = arrow_schema::Field::from(lance_field.ok_or_else(|| { - Error::internal(format!( - "Blob v2 column '{}' missing from dataset schema during compaction", - field.name() - )) - })?); - new_fields.push(Arc::new( - arrow_schema::Field::new( - field.name(), - BLOB_V2_LOGICAL_TYPE.clone(), - field.is_nullable(), - ) - .with_metadata(logical_field.metadata().clone()), - )); + new_columns.push(new_column); + new_fields.push(new_field); } let new_schema = Arc::new(arrow_schema::Schema::new_with_metadata( diff --git a/rust/lance/src/dataset/updater.rs b/rust/lance/src/dataset/updater.rs index 4314dd05b1f..194ae7c166b 100644 --- a/rust/lance/src/dataset/updater.rs +++ b/rust/lance/src/dataset/updater.rs @@ -48,6 +48,8 @@ pub struct Updater { /// The adapter to convert the logical data to physical data. schema_adapter: Option, + allow_external_blob_outside_bases: bool, + finished: bool, deletion_restorer: DeletionRestorer, @@ -97,6 +99,7 @@ impl Updater { // The schema adapter needs the data schema, not the logical schema, so it can't be // created until after the first batch is read. schema_adapter: None, + allow_external_blob_outside_bases: false, finished: false, deletion_restorer: DeletionRestorer::new(deletion_vector, legacy_batch_size), }) @@ -148,7 +151,19 @@ impl Updater { .data_storage_format .lance_file_version()?; - versions::open_update_writer(data_storage_version.into(), self.dataset(), &schema).await + versions::open_update_writer( + data_storage_version.into(), + self.dataset(), + &schema, + self.allow_external_blob_outside_bases, + ) + .await + } + + /// Allow trusted existing external blob references to pass through an update rewrite. + /// Callers must separately validate any newly supplied references before writing. + pub(super) fn allow_external_blob_outside_bases(&mut self) { + self.allow_external_blob_outside_bases = true; } /// Update one batch. diff --git a/rust/lance/src/dataset/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index 671e757eb10..da25509e372 100644 --- a/rust/lance/src/dataset/versions/mod.rs +++ b/rust/lance/src/dataset/versions/mod.rs @@ -229,6 +229,7 @@ pub async fn open_update_writer( version: ConcreteFileVersion, dataset: &Dataset, schema: &Schema, + allow_external_blob_outside_bases: bool, ) -> Result> { let external_base_resolver = match version { ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { @@ -242,7 +243,11 @@ pub async fn open_update_writer( &dataset.object_store, schema, &dataset.base, - WriterOptions::update(dataset.session.store_registry(), external_base_resolver), + WriterOptions::update( + dataset.session.store_registry(), + external_base_resolver, + allow_external_blob_outside_bases, + ), ) .await } diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 3547c088352..a30b066c7c6 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1566,10 +1566,12 @@ impl WriterOptions { pub(super) fn update( source_store_registry: Arc, external_base_resolver: Option>, + allow_external_blob_outside_bases: bool, ) -> Self { Self { add_data_dir: true, external_base_resolver, + allow_external_blob_outside_bases, source_store_registry, ..Default::default() } From 5398dd3ca6f5812b18756f107e1136aafecc7851 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:17:33 +0000 Subject: [PATCH 3/5] fix: bound blob rewrite memory --- rust/lance/src/dataset/blob.rs | 14 +++++- rust/lance/src/dataset/optimize.rs | 73 +++++++++++++++++++----------- 2 files changed, 59 insertions(+), 28 deletions(-) diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index d2e94b478d8..6e8a056c44b 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -3155,6 +3155,16 @@ async fn collect_blob_files_v2( blob_field_id: u32, descriptions: &StructArray, row_addrs: &arrow::array::PrimitiveArray, +) -> Result>> { + collect_blob_v2_descriptor_files(dataset, blob_field_id, descriptions, row_addrs.values()).await +} + +/// Resolve blob v2 descriptors to lazy handles without materializing their payloads. +pub(super) async fn collect_blob_v2_descriptor_files( + dataset: &Arc, + blob_field_id: u32, + descriptions: &StructArray, + row_addrs: &[u64], ) -> Result>> { if descriptions.len() != row_addrs.len() { return Err(Error::internal(format!( @@ -3166,7 +3176,7 @@ async fn collect_blob_files_v2( let columns = BlobV2DescriptorColumns::new(descriptions); let mut files = Vec::with_capacity(row_addrs.len()); let mut read_context = BlobV2ReadContext::new(dataset, blob_field_id); - for (selection_index, row_addr) in row_addrs.values().iter().enumerate() { + for (selection_index, row_addr) in row_addrs.iter().enumerate() { files.push( read_context .collect_file(&columns, selection_index, *row_addr) @@ -3438,7 +3448,7 @@ async fn materialize_blob_v2_list_array( Ok(Arc::new(list_array)) } -pub(super) async fn materialize_blob_v2_descriptors( +async fn materialize_blob_v2_descriptors( dataset: &Arc, blob_field_id: u32, descriptions: &StructArray, diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 7476e3376a9..4a1aa974418 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1120,8 +1120,8 @@ async fn descriptor_to_logical_blob_array( row_addrs.len() ))); } - let blob_data = if classification.data_blob_indices.is_empty() { - None + let blob_files = if classification.data_blob_indices.is_empty() { + Vec::new() } else { let indices = classification .data_blob_indices @@ -1143,20 +1143,44 @@ async fn descriptor_to_logical_blob_array( .iter() .map(|index| row_addrs[*index]) .collect::>(); - Some( - super::blob::materialize_blob_v2_descriptors( - dataset, - blob_field_id, - descriptions, - &data_row_addrs, - ) - .await?, + super::blob::collect_blob_v2_descriptor_files( + dataset, + blob_field_id, + descriptions, + &data_row_addrs, ) + .await? }; + let data_capacity = + classification + .data_blob_indices + .iter() + .try_fold(0usize, |data_capacity, row_idx| { + if descriptor.size_col.is_null(*row_idx) { + return Err(Error::internal(format!( + "Non-null blob row {} in field '{}' is missing its size", + row_idx, field_name + ))); + } + let size = usize::try_from(descriptor.size_col.value(*row_idx)).map_err(|_| { + Error::internal(format!( + "Blob size {} at row {} in field '{}' does not fit in usize", + descriptor.size_col.value(*row_idx), + row_idx, + field_name + )) + })?; + data_capacity.checked_add(size).ok_or_else(|| { + Error::internal(format!( + "Total blob size in field '{}' exceeds usize", + field_name + )) + }) + })?; + let num_rows = struct_arr.len(); - let data_values = blob_data.as_ref().map(|data| data.as_binary::()); - let mut data_builder = LargeBinaryBuilder::with_capacity(num_rows, 0); + let mut data_builder = LargeBinaryBuilder::with_capacity(num_rows, data_capacity); let mut uri_builder = StringBuilder::with_capacity(num_rows, 0); let mut out_position_builder = PrimitiveBuilder::::with_capacity(num_rows); let mut out_size_builder = PrimitiveBuilder::::with_capacity(num_rows); @@ -1199,21 +1223,18 @@ async fn descriptor_to_logical_blob_array( } } RowClass::DataBlob => { - let data_values = data_values.as_ref().ok_or_else(|| { - Error::internal(format!( - "Non-null blob row {} in field '{}' produced no materialized data", - i, field_name - )) - })?; - if data_values.is_null(blob_file_idx) { - return Err(Error::internal(format!( - "Non-null blob row {} in field '{}' resolved to null", - i, field_name - ))); - } - let data = data_values.value(blob_file_idx); + let blob_file = blob_files + .get(blob_file_idx) + .and_then(Option::as_ref) + .ok_or_else(|| { + Error::internal(format!( + "Non-null blob row {} in field '{}' resolved to null", + i, field_name + )) + })?; + let data = blob_file.read().await?; blob_file_idx += 1; - data_builder.append_value(data); + data_builder.append_value(data.as_ref()); uri_builder.append_null(); out_position_builder.append_null(); out_size_builder.append_null(); From 367303436de6ccb1b5a10eec10692fb817fbaf8a Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:38:17 +0000 Subject: [PATCH 4/5] fix: align blob v2 rewrite schema and traversal --- python/python/tests/test_optimize.py | 39 ++ rust/lance/src/dataset/blob.rs | 18 + rust/lance/src/dataset/optimize.rs | 623 ++++++++++++++++++--------- 3 files changed, 475 insertions(+), 205 deletions(-) diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 82cfb4e33a4..b6815b32da2 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +import json import pickle import random import re @@ -88,6 +89,44 @@ def test_blob_compaction(tmp_path: Path): assert contents == blobs +def test_blob_compaction_with_nested_json_sibling(tmp_path: Path): + dataset_uri = tmp_path / "nested_blob_json" + info_fields = [lance.blob_field("blob"), pa.field("meta", pa.json_())] + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("info", pa.struct(info_fields)), + ] + ) + for index, row_id in enumerate([1, 2]): + info = pa.StructArray.from_arrays( + [ + lance.blob_array([f"blob-{row_id}".encode()]), + pa.array([json.dumps({"row": row_id})], type=pa.json_()), + ], + fields=info_fields, + ) + lance.write_dataset( + pa.Table.from_arrays([pa.array([row_id]), info], schema=schema), + dataset_uri, + mode="create" if index == 0 else "append", + data_storage_version="2.2", + ) + + dataset = lance.dataset(dataset_uri) + dataset.optimize.compact_files(num_threads=1) + + assert len(dataset.get_fragments()) == 1 + assert [data for _, data in dataset.read_blobs("info.blob", indices=[0, 1])] == [ + b"blob-1", + b"blob-2", + ] + assert [ + json.loads(value) + for value in dataset.to_table(columns=["info.meta"])["info.meta"].to_pylist() + ] == [{"row": 1}, {"row": 2}] + + @pytest.mark.parametrize("storage_version", ["2.0", "2.1", "2.2"]) def test_blob_compaction_preserves_null_empty_and_read_parity( tmp_path: Path, storage_version: str diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 6e8a056c44b..078466ed76f 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -188,6 +188,21 @@ impl ExternalBaseResolver { } } +fn arrow_field_contains_blob_v2(field: &ArrowField) -> bool { + if field.is_blob_v2() { + return true; + } + match field.data_type() { + ArrowDataType::Struct(children) => children + .iter() + .any(|child| arrow_field_contains_blob_v2(child)), + ArrowDataType::List(child) | ArrowDataType::LargeList(child) => { + arrow_field_contains_blob_v2(child) + } + _ => false, + } +} + fn collect_external_blob_uris( field: &ArrowField, array: &ArrayRef, @@ -195,6 +210,9 @@ fn collect_external_blob_uris( field_path: &str, external_uris: &mut Vec<(String, String)>, ) -> Result<()> { + if !arrow_field_contains_blob_v2(field) { + return Ok(()); + } if array.len() != selected_rows.len() { return Err(Error::internal(format!( "Blob field '{}' row count {} did not match selection length {}", diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 4a1aa974418..1167df59120 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -107,7 +107,9 @@ use arrow_array::{ Array, ArrayRef, GenericListArray, OffsetSizeTrait, RecordBatch, StructArray, UInt32Array, }; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; -use arrow_schema::{DataType as ArrowDataType, Field as ArrowField}; +use arrow_schema::{ + DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef, +}; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::future::BoxFuture; @@ -1256,133 +1258,296 @@ async fn descriptor_to_logical_blob_array( fn transformed_arrow_field(field: &LanceField, data_type: ArrowDataType) -> Arc { let arrow_field = ArrowField::from(field); + arrow_field_with_data_type(&arrow_field, data_type) +} + +fn arrow_field_with_data_type(field: &ArrowField, data_type: ArrowDataType) -> Arc { Arc::new( - ArrowField::new(arrow_field.name(), data_type, arrow_field.is_nullable()) - .with_metadata(arrow_field.metadata().clone()), + ArrowField::new(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()), ) } -fn transform_blob_v2_array<'a>( - dataset: &'a Arc, - field: &'a LanceField, - array: ArrayRef, - row_addrs: Arc<[u64]>, -) -> BoxFuture<'a, Result<(ArrayRef, Arc)>> { - async move { +fn field_contains_blob_v2(field: &LanceField) -> bool { + field.is_blob_v2() || field.children.iter().any(field_contains_blob_v2) +} + +enum BlobV2FieldRewritePlan { + Passthrough { + output_field: Arc, + }, + Blob { + field_id: u32, + field_name: String, + output_field: Arc, + }, + Struct { + field_name: String, + output_field: Arc, + children: Vec, + }, + List { + field_name: String, + output_field: Arc, + child: Box, + }, + LargeList { + field_name: String, + output_field: Arc, + child: Box, + }, +} + +impl BlobV2FieldRewritePlan { + fn passthrough(field: &ArrowField) -> Self { + Self::Passthrough { + output_field: Arc::new(field.clone()), + } + } + + fn try_new(field: &LanceField, input_field: &ArrowField) -> Result { + if !field_contains_blob_v2(field) { + return Ok(Self::passthrough(input_field)); + } + if field.is_blob_v2() { - let struct_arr = array.as_struct(); - return match BlobV2Layout::classify(struct_arr.fields()) { - Some(BlobV2Layout::Logical) => { - let output_field = transformed_arrow_field(field, array.data_type().clone()); - Ok((array, output_field)) - } + let field_id = u32::try_from(field.id).map_err(|_| { + Error::internal(format!( + "Blob v2 field id {} for '{}' does not fit in u32", + field.id, field.name + )) + })?; + let ArrowDataType::Struct(input_children) = input_field.data_type() else { + return Err(Error::invalid_input(format!( + "Blob v2 field '{}' has non-struct input type {:?}", + field.name, + input_field.data_type() + ))); + }; + let output_field = match BlobV2Layout::classify(input_children) { + Some(BlobV2Layout::Logical) => Arc::new(input_field.clone()), Some(BlobV2Layout::Descriptor) => { - let descriptor = BlobV2Descriptor::try_from_struct(struct_arr, &field.name)?; - let classification = classify_rows(struct_arr, &descriptor, &field.name)?; - let logical = descriptor_to_logical_blob_array( - dataset, - u32::try_from(field.id).map_err(|_| { - Error::internal(format!( - "Blob v2 field id {} for '{}' does not fit in u32", - field.id, field.name - )) - })?, - struct_arr, - &descriptor, - &classification, - row_addrs.as_ref(), - &field.name, - ) - .await?; - Ok(( - Arc::new(logical) as ArrayRef, - transformed_arrow_field(field, BLOB_V2_LOGICAL_TYPE.clone()), - )) + transformed_arrow_field(field, BLOB_V2_LOGICAL_TYPE.clone()) + } + Some(actual) => { + return Err(Error::invalid_input(format!( + "Blob v2 field '{}' has {actual} input layout; expected logical or descriptor layout during rewrite", + field.name + ))); + } + None => { + return Err(Error::invalid_input(format!( + "Blob v2 field '{}' has unrecognized input layout {:?}; expected logical or descriptor layout during rewrite", + field.name, input_children + ))); } - Some(actual) => Err(Error::invalid_input(format!( - "Blob v2 field '{}' has {actual} layout; expected logical or descriptor layout during rewrite", - field.name - ))), - None => Err(Error::invalid_input(format!( - "Blob v2 field '{}' has unrecognized layout {:?}; expected logical or descriptor layout during rewrite", - field.name, - struct_arr.fields() - ))), }; + return Ok(Self::Blob { + field_id, + field_name: field.name.clone(), + output_field, + }); } - match field.data_type() { - ArrowDataType::Struct(_) => { - let struct_arr = array.as_struct().normalize_slicing()?; - let parent_nulls = struct_arr.nulls().cloned(); - let struct_arr = struct_arr.pushdown_nulls()?; - if field.children.len() != struct_arr.num_columns() { + match (field.data_type(), input_field.data_type()) { + (ArrowDataType::Struct(_), ArrowDataType::Struct(input_children)) => { + if field.children.len() != input_children.len() { return Err(Error::internal(format!( - "Struct field '{}' expected {} children during blob rewrite, got {}", + "Struct field '{}' expected {} children in blob rewrite plan, got {}", field.name, field.children.len(), - struct_arr.num_columns() + input_children.len() ))); } - let mut child_arrays = Vec::with_capacity(field.children.len()); - let mut child_fields = Vec::with_capacity(field.children.len()); - for (child_field, child_array) in - field.children.iter().zip(struct_arr.columns()) - { - let (child_array, child_field) = transform_blob_v2_array( - dataset, - child_field, - child_array.clone(), - row_addrs.clone(), - ) - .await?; - child_arrays.push(child_array); - child_fields.push(child_field); - } - let struct_array = StructArray::try_new( - child_fields.clone().into(), - child_arrays, - parent_nulls, - )?; - let output_field = transformed_arrow_field( - field, - ArrowDataType::Struct(child_fields.into()), - ); - Ok((Arc::new(struct_array), output_field)) + let children = field + .children + .iter() + .zip(input_children.iter()) + .map(|(child, input_child)| Self::try_new(child, input_child)) + .collect::>>()?; + let output_children = children + .iter() + .map(|child| child.output_field().clone()) + .collect::>(); + Ok(Self::Struct { + field_name: field.name.clone(), + output_field: arrow_field_with_data_type( + input_field, + ArrowDataType::Struct(output_children.into()), + ), + children, + }) } - ArrowDataType::List(_) => { - transform_blob_v2_list_array::( - dataset, - field, - array.as_list::(), - row_addrs, - ) - .await + (ArrowDataType::List(_), ArrowDataType::List(input_child)) => { + let child = field.children.first().ok_or_else(|| { + Error::internal(format!( + "List field '{}' is missing its child in blob rewrite plan", + field.name + )) + })?; + let child = Box::new(Self::try_new(child, input_child)?); + Ok(Self::List { + field_name: field.name.clone(), + output_field: arrow_field_with_data_type( + input_field, + ArrowDataType::List(child.output_field().clone()), + ), + child, + }) } - ArrowDataType::LargeList(_) => { - transform_blob_v2_list_array::( - dataset, - field, - array.as_list::(), - row_addrs, - ) - .await + (ArrowDataType::LargeList(_), ArrowDataType::LargeList(input_child)) => { + let child = field.children.first().ok_or_else(|| { + Error::internal(format!( + "Large list field '{}' is missing its child in blob rewrite plan", + field.name + )) + })?; + let child = Box::new(Self::try_new(child, input_child)?); + Ok(Self::LargeList { + field_name: field.name.clone(), + output_field: arrow_field_with_data_type( + input_field, + ArrowDataType::LargeList(child.output_field().clone()), + ), + child, + }) } - _ => { - let output_field = transformed_arrow_field(field, array.data_type().clone()); - Ok((array, output_field)) + (logical_type, input_type) => Err(Error::invalid_input(format!( + "Field '{}' contains blob v2 descendants but logical type {:?} and input type {:?} do not form a supported rewrite container", + field.name, logical_type, input_type + ))), + } + } + + fn output_field(&self) -> &Arc { + match self { + Self::Passthrough { output_field } + | Self::Blob { output_field, .. } + | Self::Struct { output_field, .. } + | Self::List { output_field, .. } + | Self::LargeList { output_field, .. } => output_field, + } + } + + fn transform<'a>( + &'a self, + dataset: &'a Arc, + array: ArrayRef, + row_addrs: Arc<[u64]>, + ) -> BoxFuture<'a, Result> { + async move { + match self { + Self::Passthrough { .. } => Ok(array), + Self::Blob { + field_id, + field_name, + .. + } => { + let struct_arr = array.as_struct(); + match BlobV2Layout::classify(struct_arr.fields()) { + Some(BlobV2Layout::Logical) => Ok(array), + Some(BlobV2Layout::Descriptor) => { + let descriptor = + BlobV2Descriptor::try_from_struct(struct_arr, field_name)?; + let classification = + classify_rows(struct_arr, &descriptor, field_name)?; + let logical = descriptor_to_logical_blob_array( + dataset, + *field_id, + struct_arr, + &descriptor, + &classification, + row_addrs.as_ref(), + field_name, + ) + .await?; + Ok(Arc::new(logical) as ArrayRef) + } + Some(actual) => Err(Error::invalid_input(format!( + "Blob v2 field '{}' has {actual} layout; expected logical or descriptor layout during rewrite", + field_name + ))), + None => Err(Error::invalid_input(format!( + "Blob v2 field '{}' has unrecognized layout {:?}; expected logical or descriptor layout during rewrite", + field_name, + struct_arr.fields() + ))), + } + } + Self::Struct { + field_name, + output_field, + children, + } => { + let struct_arr = array.as_struct().normalize_slicing()?; + let parent_nulls = struct_arr.nulls().cloned(); + let struct_arr = struct_arr.pushdown_nulls()?; + if children.len() != struct_arr.num_columns() { + return Err(Error::internal(format!( + "Struct field '{}' expected {} children during blob rewrite, got {}", + field_name, + children.len(), + struct_arr.num_columns() + ))); + } + let mut child_arrays = Vec::with_capacity(children.len()); + for (child, child_array) in children.iter().zip(struct_arr.columns()) { + child_arrays.push( + child + .transform(dataset, child_array.clone(), row_addrs.clone()) + .await?, + ); + } + let ArrowDataType::Struct(output_fields) = output_field.data_type() else { + return Err(Error::internal(format!( + "Struct field '{}' rewrite plan has non-struct output type {:?}", + field_name, + output_field.data_type() + ))); + }; + Ok(Arc::new(StructArray::try_new( + output_fields.clone(), + child_arrays, + parent_nulls, + )?) as ArrayRef) + } + Self::List { + field_name, child, .. + } => { + transform_blob_v2_list_array::( + dataset, + field_name, + child, + array.as_list::(), + row_addrs, + ) + .await + } + Self::LargeList { + field_name, child, .. + } => { + transform_blob_v2_list_array::( + dataset, + field_name, + child, + array.as_list::(), + row_addrs, + ) + .await + } } } + .boxed() } - .boxed() } async fn transform_blob_v2_list_array( dataset: &Arc, - field: &LanceField, + field_name: &str, + child: &BlobV2FieldRewritePlan, list_array: &GenericListArray, row_addrs: Arc<[u64]>, -) -> Result<(ArrayRef, Arc)> { +) -> Result { let list_array = if list_array.null_count() > 0 { list_array.filter_garbage_nulls() } else { @@ -1394,7 +1559,7 @@ async fn transform_blob_v2_list_array( if values_end < values_start { return Err(Error::internal(format!( "List field '{}' has invalid offsets during blob rewrite", - field.name + field_name ))); } @@ -1408,13 +1573,13 @@ async fn transform_blob_v2_list_array( if end < start { return Err(Error::internal(format!( "List field '{}' has decreasing offsets during blob rewrite", - field.name + field_name ))); } let row_addr = row_addrs.get(row_idx).copied().ok_or_else(|| { Error::internal(format!( "List field '{}' row address count {} did not match row count {}", - field.name, + field_name, row_addrs.len(), list_array.len() )) @@ -1423,90 +1588,109 @@ async fn transform_blob_v2_list_array( normalized_offsets.push(O::usize_as(end - values_start)); } - let child = field.children.first().ok_or_else(|| { - Error::internal(format!( - "List field '{}' missing child during blob rewrite", - field.name - )) - })?; let values = list_array.values().slice(values_start, values_len); - let (values, child_field) = - transform_blob_v2_array(dataset, child, values, Arc::<[u64]>::from(child_row_addrs)) - .await?; + let values = child + .transform(dataset, values, Arc::<[u64]>::from(child_row_addrs)) + .await?; let list_array = GenericListArray::::try_new( - child_field, + child.output_field().clone(), OffsetBuffer::new(ScalarBuffer::from(normalized_offsets)), values, list_array.nulls().cloned(), )?; - let output_field = transformed_arrow_field(field, list_array.data_type().clone()); - Ok((Arc::new(list_array), output_field)) + Ok(Arc::new(list_array)) } -pub(crate) async fn transform_blob_v2_batch( - dataset: &Arc, - schema: &lance_core::datatypes::Schema, - batch: RecordBatch, - keep_row_addr: bool, -) -> Result { - let row_addr_idx = batch - .schema() - .column_with_name(lance_core::ROW_ADDR) - .ok_or_else(|| { - Error::internal(format!( - "_rowaddr column missing from batch for blob v2 compaction, columns: {:?}", - batch - .schema() - .fields() - .iter() - .map(|f| f.name()) - .collect::>() - )) - })? - .0; - let row_addrs: Arc<[u64]> = batch - .column(row_addr_idx) - .as_primitive::() - .values() - .iter() - .copied() - .collect::>() - .into(); - - let mut new_columns: Vec> = Vec::new(); - let mut new_fields: Vec> = Vec::new(); +struct BlobV2BatchRewritePlan { + row_addr_idx: usize, + columns: Vec<(usize, BlobV2FieldRewritePlan)>, + output_schema: SchemaRef, +} - let batch_schema = batch.schema(); - for (col_idx, field) in batch_schema.fields().iter().enumerate() { - if field.name() == lance_core::ROW_ADDR && !keep_row_addr { - continue; +impl BlobV2BatchRewritePlan { + fn try_new( + schema: &lance_core::datatypes::Schema, + input_schema: &ArrowSchema, + keep_row_addr: bool, + ) -> Result { + let row_addr_idx = input_schema + .column_with_name(lance_core::ROW_ADDR) + .ok_or_else(|| { + Error::internal(format!( + "_rowaddr column missing from batch for blob v2 compaction, columns: {:?}", + input_schema + .fields() + .iter() + .map(|f| f.name()) + .collect::>() + )) + })? + .0; + let mut columns = Vec::with_capacity(input_schema.fields().len()); + let mut output_fields = Vec::with_capacity(input_schema.fields().len()); + for (column_idx, input_field) in input_schema.fields().iter().enumerate() { + if input_field.name() == lance_core::ROW_ADDR && !keep_row_addr { + continue; + } + let field_plan = if let Some(field) = schema.field(input_field.name()) { + BlobV2FieldRewritePlan::try_new(field, input_field)? + } else { + BlobV2FieldRewritePlan::passthrough(input_field) + }; + output_fields.push(field_plan.output_field().as_ref().clone()); + columns.push((column_idx, field_plan)); } - let Some(lance_field) = schema.field(field.name()) else { - new_columns.push(batch.column(col_idx).clone()); - new_fields.push(field.clone()); - continue; - }; - let (new_column, new_field) = transform_blob_v2_array( - dataset, - lance_field, - batch.column(col_idx).clone(), - row_addrs.clone(), - ) - .await?; - new_columns.push(new_column); - new_fields.push(new_field); + Ok(Self { + row_addr_idx, + columns, + output_schema: Arc::new(ArrowSchema::new_with_metadata( + output_fields, + input_schema.metadata().clone(), + )), + }) } - let new_schema = Arc::new(arrow_schema::Schema::new_with_metadata( - new_fields + async fn transform_batch( + &self, + dataset: &Arc, + batch: RecordBatch, + ) -> Result { + let row_addrs: Arc<[u64]> = batch + .column(self.row_addr_idx) + .as_primitive::() + .values() .iter() - .map(|f| f.as_ref().clone()) - .collect::>(), - batch_schema.metadata().clone(), - )); + .copied() + .collect::>() + .into(); + let mut output_columns = Vec::with_capacity(self.columns.len()); + for (column_idx, field_plan) in &self.columns { + output_columns.push( + field_plan + .transform( + dataset, + batch.column(*column_idx).clone(), + row_addrs.clone(), + ) + .await?, + ); + } + Ok(RecordBatch::try_new( + self.output_schema.clone(), + output_columns, + )?) + } +} - Ok(RecordBatch::try_new(new_schema, new_columns)?) +pub(crate) async fn transform_blob_v2_batch( + dataset: &Arc, + schema: &lance_core::datatypes::Schema, + batch: RecordBatch, + keep_row_addr: bool, +) -> Result { + let plan = BlobV2BatchRewritePlan::try_new(schema, batch.schema().as_ref(), keep_row_addr)?; + plan.transform_batch(dataset, batch).await } /// Build a scan reader for rewrite and optionally capture row IDs. @@ -1902,46 +2086,23 @@ async fn rewrite_files( if has_blob_v2_columns { let dataset_arc = Arc::new(dataset.as_ref().clone()); - let dataset_schema = dataset.schema().clone(); + let rewrite_plan = Arc::new(BlobV2BatchRewritePlan::try_new( + dataset.schema(), + schema.as_ref(), + false, + )?); + let transformed_schema = rewrite_plan.output_schema.clone(); let transformed = reader_with_progress.then(move |batch_result| { let dataset = dataset_arc.clone(); - let schema = dataset_schema.clone(); + let rewrite_plan = rewrite_plan.clone(); async move { let batch = batch_result?; - transform_blob_v2_batch(&dataset, &schema, batch, false) + rewrite_plan + .transform_batch(&dataset, batch) .await .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e))) } }); - let transformed_schema = { - let mut fields: Vec> = Vec::new(); - for field in schema.fields().iter() { - if field.name() == lance_core::ROW_ADDR { - continue; - } - let lance_field = dataset.schema().field(field.name()); - if let Some(lance_field) = lance_field.filter(|f| f.is_blob_v2()) { - let logical_field = arrow_schema::Field::from(lance_field); - fields.push(Arc::new( - arrow_schema::Field::new( - field.name(), - BLOB_V2_LOGICAL_TYPE.clone(), - field.is_nullable(), - ) - .with_metadata(logical_field.metadata().clone()), - )); - } else { - fields.push(field.clone()); - } - } - Arc::new(arrow_schema::Schema::new_with_metadata( - fields - .iter() - .map(|f| f.as_ref().clone()) - .collect::>(), - schema.metadata().clone(), - )) - }; reader = Some(Box::pin(RecordBatchStreamAdapter::new( transformed_schema, transformed, @@ -7812,6 +7973,58 @@ mod tests { assert_eq!(after, expected); } + #[test] + fn test_blob_v2_rewrite_plan_skips_non_blob_list_subtree() { + let items_field = Field::new( + "items", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ); + let info_field = Field::new( + "info", + DataType::Struct( + vec![ + Arc::new(items_field), + Arc::new(crate::blob_field("blob", true)), + ] + .into(), + ), + true, + ); + let logical_arrow_schema = Schema::new(vec![info_field]); + let logical_schema = + lance_core::datatypes::Schema::try_from(&logical_arrow_schema).unwrap(); + let mut input_schema = logical_schema.clone(); + input_schema.fields[0].unload_blobs_recursive(); + let input_field = Field::from(&input_schema.fields[0]); + + let plan = + BlobV2FieldRewritePlan::try_new(&logical_schema.fields[0], &input_field).unwrap(); + let BlobV2FieldRewritePlan::Struct { children, .. } = &plan else { + panic!("nested blob field should produce a struct rewrite plan"); + }; + assert!(matches!( + children[0], + BlobV2FieldRewritePlan::Passthrough { .. } + )); + assert!(matches!(children[1], BlobV2FieldRewritePlan::Blob { .. })); + + let DataType::Struct(output_children) = plan.output_field().data_type() else { + panic!("nested blob rewrite plan should declare a struct output"); + }; + let DataType::Struct(input_children) = input_field.data_type() else { + panic!("nested blob input should be a struct"); + }; + assert_eq!(output_children[0], input_children[0]); + let DataType::Struct(blob_children) = output_children[1].data_type() else { + panic!("blob rewrite plan should declare a struct blob output"); + }; + assert_eq!( + BlobV2Layout::classify(blob_children), + Some(BlobV2Layout::Logical) + ); + } + #[tokio::test] async fn test_compact_blob_v1_preserves_null_empty_and_payload_order() { let test_dir = TempStrDir::default(); From 942bfa1bbd3f8bbec834d3cdc76e3061b4ef157c Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:10:50 +0000 Subject: [PATCH 5/5] fix: reuse blob rewrite plan in merge insert --- python/python/tests/test_dataset.py | 41 +++++++ rust/lance/src/dataset/optimize.rs | 12 +- .../dataset/write/merge_insert/exec/write.rs | 107 +++--------------- 3 files changed, 64 insertions(+), 96 deletions(-) diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 67c74663f2d..f165056fdb3 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -2656,6 +2656,47 @@ def test_merge_insert_subcols(tmp_path: Path): assert dataset.to_table().sort_by("a") == expected +@pytest.mark.parametrize("container", ["struct", "list"]) +def test_merge_insert_subcols_preserves_nested_blob(tmp_path: Path, container: str): + blob_field = lance.blob_field("blob") + blob_values = lance.blob_array([b"one", b"two"]) + if container == "struct": + nested_values = pa.StructArray.from_arrays( + [blob_values], + fields=[blob_field], + ) + expected_nested = [{"blob": b"one"}, {"blob": b"two"}] + else: + nested_values = pa.ListArray.from_arrays( + pa.array([0, 1, 2], type=pa.int32()), + blob_values, + type=pa.list_(blob_field), + ) + expected_nested = [[b"one"], [b"two"]] + + dataset_uri = tmp_path / f"partial_nested_blob_{container}" + dataset = lance.write_dataset( + pa.table( + { + "id": pa.array([1, 2]), + "nested": nested_values, + "other": pa.array([10, 20]), + } + ), + dataset_uri, + data_storage_version="2.2", + ) + source = pa.table({"id": pa.array([2]), "other": pa.array([200])}) + + dataset.merge_insert("id").when_matched_update_all().execute(source) + + result = ( + lance.dataset(dataset_uri).to_table(blob_handling="all_binary").sort_by("id") + ) + assert result["other"].to_pylist() == [10, 200] + assert result["nested"].to_pylist() == expected_nested + + def test_merge_insert_full_fragment_rewrite_json_e2e(tmp_path: Path): """End-to-end test: merge_insert with JSON columns where ALL rows are updated. diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 1167df59120..f9abba1557e 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1601,14 +1601,14 @@ async fn transform_blob_v2_list_array( Ok(Arc::new(list_array)) } -struct BlobV2BatchRewritePlan { +pub(crate) struct BlobV2BatchRewritePlan { row_addr_idx: usize, columns: Vec<(usize, BlobV2FieldRewritePlan)>, output_schema: SchemaRef, } impl BlobV2BatchRewritePlan { - fn try_new( + pub(crate) fn try_new( schema: &lance_core::datatypes::Schema, input_schema: &ArrowSchema, keep_row_addr: bool, @@ -1617,7 +1617,7 @@ impl BlobV2BatchRewritePlan { .column_with_name(lance_core::ROW_ADDR) .ok_or_else(|| { Error::internal(format!( - "_rowaddr column missing from batch for blob v2 compaction, columns: {:?}", + "_rowaddr column missing from batch for blob v2 rewrite, columns: {:?}", input_schema .fields() .iter() @@ -1651,7 +1651,11 @@ impl BlobV2BatchRewritePlan { }) } - async fn transform_batch( + pub(crate) fn output_schema(&self) -> &SchemaRef { + &self.output_schema + } + + pub(crate) async fn transform_batch( &self, dataset: &Arc, batch: RecordBatch, diff --git a/rust/lance/src/dataset/write/merge_insert/exec/write.rs b/rust/lance/src/dataset/write/merge_insert/exec/write.rs index 3bf3e1237c8..55f5ce4070f 100644 --- a/rust/lance/src/dataset/write/merge_insert/exec/write.rs +++ b/rust/lance/src/dataset/write/merge_insert/exec/write.rs @@ -22,10 +22,7 @@ use datafusion::{ }; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use futures::{StreamExt, stream}; -use lance_core::{ - Error, ROW_ADDR, ROW_ID, - datatypes::{BLOB_V2_LOGICAL_TYPE, BlobV2Layout}, -}; +use lance_core::{Error, ROW_ADDR, ROW_ID}; use lance_table::format::RowIdMeta; use roaring::RoaringTreemap; @@ -55,60 +52,6 @@ use crate::{ use super::apply_deletions; -fn descriptor_to_logical_blob_schema( - input_schema: arrow_schema::SchemaRef, - dataset_schema: &lance_core::datatypes::Schema, -) -> lance_core::Result { - let fields = input_schema - .fields() - .iter() - .map(|field| -> lance_core::Result<_> { - let Some(dataset_field) = dataset_schema - .field(field.name()) - .filter(|dataset_field| dataset_field.is_blob_v2()) - else { - return Ok(field.clone()); - }; - let arrow_schema::DataType::Struct(fields) = field.data_type() else { - return Err(Error::invalid_input(format!( - "Blob v2 merge input '{}' has non-struct type {}; expected logical or descriptor layout", - field.name(), - field.data_type() - ))); - }; - match BlobV2Layout::classify(fields) { - Some(BlobV2Layout::Logical) => Ok(field.clone()), - Some(BlobV2Layout::Descriptor) => { - let logical_field = arrow_schema::Field::from(dataset_field); - Ok(Arc::new( - arrow_schema::Field::new( - field.name(), - BLOB_V2_LOGICAL_TYPE.clone(), - field.is_nullable(), - ) - .with_metadata(logical_field.metadata().clone()), - )) - } - Some(actual) => Err(Error::invalid_input(format!( - "Blob v2 merge input '{}' has {actual} layout; expected logical or descriptor layout", - field.name() - ))), - None => Err(Error::invalid_input(format!( - "Blob v2 merge input '{}' has unrecognized layout {fields:?}; expected logical or descriptor layout", - field.name() - ))), - } - }) - .collect::>>()?; - Ok(Arc::new(arrow_schema::Schema::new_with_metadata( - fields - .iter() - .map(|field| field.as_ref().clone()) - .collect::>(), - input_schema.metadata().clone(), - ))) -} - /// Shared state for merge insert operations to simplify lock management struct MergeState { /// Row addresses that need to be deleted, due to a row update or delete action @@ -946,24 +889,25 @@ impl ExecutionPlan for FullSchemaMergeInsertExec { .any(|field| field.is_blob_v2()); let input_stream = if has_blob_v2_columns { let input_schema = input_stream.schema(); - let output_schema = - descriptor_to_logical_blob_schema(input_schema, self.dataset.schema()) - .map_err(|error| DataFusionError::External(Box::new(error)))?; + let rewrite_plan = Arc::new( + crate::dataset::optimize::BlobV2BatchRewritePlan::try_new( + self.dataset.schema(), + input_schema.as_ref(), + true, + ) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + let output_schema = rewrite_plan.output_schema().clone(); let dataset = self.dataset.clone(); - let dataset_schema = self.dataset.schema().clone(); let transformed = input_stream.then(move |batch_result| { let dataset = dataset.clone(); - let dataset_schema = dataset_schema.clone(); + let rewrite_plan = rewrite_plan.clone(); async move { let batch = batch_result?; - crate::dataset::optimize::transform_blob_v2_batch( - &dataset, - &dataset_schema, - batch, - true, - ) - .await - .map_err(|error| DataFusionError::External(Box::new(error))) + rewrite_plan + .transform_batch(&dataset, batch) + .await + .map_err(|error| DataFusionError::External(Box::new(error))) } }); Box::pin(RecordBatchStreamAdapter::new(output_schema, transformed)) @@ -1171,27 +1115,6 @@ mod tests { use super::*; use arrow_array::UInt64Array; - #[test] - fn test_descriptor_to_logical_blob_schema_rejects_prepared_layout() { - let logical_field = crate::blob::blob_field("blob", true); - let dataset_schema = - lance_core::datatypes::Schema::try_from(&Schema::new(vec![logical_field.clone()])) - .unwrap(); - let prepared_field = arrow_schema::Field::new( - "blob", - lance_core::datatypes::BLOB_V2_PREPARED_TYPE.clone(), - true, - ) - .with_metadata(logical_field.metadata().clone()); - let input_schema = Arc::new(Schema::new(vec![prepared_field])); - - let error = descriptor_to_logical_blob_schema(input_schema, &dataset_schema).unwrap_err(); - assert!(matches!(error, Error::InvalidInput { .. })); - assert!(error.to_string().contains( - "Blob v2 merge input 'blob' has prepared layout; expected logical or descriptor layout" - )); - } - #[test] fn test_merge_state_duplicate_rowid_detection_fail() { let metrics = MergeInsertMetrics::new(&ExecutionPlanMetricsSet::new(), 0);