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/python/python/tests/test_fragment.py b/python/python/tests/test_fragment.py index 11276a2213d..deb48c57076 100644 --- a/python/python/tests/test_fragment.py +++ b/python/python/tests/test_fragment.py @@ -617,6 +617,139 @@ 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_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/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 b59250d670f..078466ed76f 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,158 @@ 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, + selected_rows: &[bool], + 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 {}", + 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, @@ -3021,6 +3173,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!( @@ -3032,7 +3194,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) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 115d024c004..b1a419a0129 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,13 +1921,50 @@ impl FileFragment { if !read_columns.iter().any(|n| n.as_str() == ROW_ADDR) { read_columns.push(ROW_ADDR.to_string()); } + 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::>(); + 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?; + 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. @@ -1937,6 +1986,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,8 +2012,12 @@ impl FileFragment { } } let updated_batch = joiner - .collect_with_fallback(batch, index_column, self.dataset()) + .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?; } @@ -5548,7 +5612,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 +5870,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/optimize.rs b/rust/lance/src/dataset/optimize.rs index 120a61c2fe8..f9abba1557e 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -102,17 +102,23 @@ 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, Schema as ArrowSchema, SchemaRef, +}; 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 +1060,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 +1090,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,20 +1107,82 @@ 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() { + 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_files = if classification.data_blob_indices.is_empty() { Vec::new() } 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::>(); + super::blob::collect_blob_v2_descriptor_files( + dataset, + blob_field_id, + descriptions, + &data_row_addrs, + ) + .await? }; - let mut data_builder = LargeBinaryBuilder::with_capacity(num_rows, 0); + 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 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); @@ -1140,7 +1207,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,12 +1225,15 @@ async fn descriptor_to_logical_blob_array( } } RowClass::DataBlob => { - let blob_file = blob_files[blob_file_idx].as_ref().ok_or_else(|| { - Error::internal(format!( - "Non-null blob row {} in column '{}' resolved to null", - i, column_name - )) - })?; + 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.as_ref()); @@ -1182,128 +1252,449 @@ 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(), )?) } -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 = batch.column(row_addr_idx).as_primitive::(); +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(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()), + ) +} + +fn field_contains_blob_v2(field: &LanceField) -> bool { + field.is_blob_v2() || field.children.iter().any(field_contains_blob_v2) +} - let mut new_columns: Vec> = Vec::new(); - let mut new_fields: Vec> = Vec::new(); +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, + }, +} - 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 BlobV2FieldRewritePlan { + fn passthrough(field: &ArrowField) -> Self { + Self::Passthrough { + output_field: Arc::new(field.clone()), } + } - 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 { - new_columns.push(batch.column(col_idx).clone()); - new_fields.push(field.clone()); - continue; + fn try_new(field: &LanceField, input_field: &ArrowField) -> Result { + if !field_contains_blob_v2(field) { + return Ok(Self::passthrough(input_field)); } - let struct_arr = batch - .column(col_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { + if field.is_blob_v2() { + let field_id = u32::try_from(field.id).map_err(|_| { Error::internal(format!( - "Blob v2 column '{}' expected StructArray, got {:?}", - field.name(), - batch.column(col_idx).data_type() + "Blob v2 field id {} for '{}' does not fit in u32", + field.id, field.name )) })?; - - 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) => { + let ArrowDataType::Struct(input_children) = input_field.data_type() else { return Err(Error::invalid_input(format!( - "Blob v2 column '{}' has {actual} layout; expected logical or descriptor layout during compaction", - field.name() + "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) => { + 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 + ))); + } + }; + return Ok(Self::Blob { + field_id, + field_name: field.name.clone(), + output_field, + }); + } + + 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 in blob rewrite plan, got {}", + field.name, + field.children.len(), + input_children.len() + ))); + } + 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, + }) } - 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() - ))); + (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(_), 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, + }) } + (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 + ))), } + } - 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(); + 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, + } + } - let new_struct = descriptor_to_logical_blob_array( - dataset, - &descriptor, - &classification, - column_name, - num_rows, - struct_arr.nulls().cloned(), - ) - .await?; + 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() + } +} - new_columns.push(Arc::new(new_struct)); - let logical_field = arrow_schema::Field::from(lance_field.ok_or_else(|| { +async fn transform_blob_v2_list_array( + dataset: &Arc, + field_name: &str, + child: &BlobV2FieldRewritePlan, + list_array: &GenericListArray, + row_addrs: Arc<[u64]>, +) -> Result { + 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!( - "Blob v2 column '{}' missing from dataset schema during compaction", - field.name() + "List field '{}' row address count {} did not match row count {}", + field_name, + row_addrs.len(), + list_array.len() )) - })?); - 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()), - )); + })?; + child_row_addrs.extend(std::iter::repeat_n(row_addr, end - start)); + normalized_offsets.push(O::usize_as(end - values_start)); } - let new_schema = Arc::new(arrow_schema::Schema::new_with_metadata( - new_fields + let values = list_array.values().slice(values_start, values_len); + let values = child + .transform(dataset, values, Arc::<[u64]>::from(child_row_addrs)) + .await?; + let list_array = GenericListArray::::try_new( + child.output_field().clone(), + OffsetBuffer::new(ScalarBuffer::from(normalized_offsets)), + values, + list_array.nulls().cloned(), + )?; + Ok(Arc::new(list_array)) +} + +pub(crate) struct BlobV2BatchRewritePlan { + row_addr_idx: usize, + columns: Vec<(usize, BlobV2FieldRewritePlan)>, + output_schema: SchemaRef, +} + +impl BlobV2BatchRewritePlan { + pub(crate) 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 rewrite, 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)); + } + + Ok(Self { + row_addr_idx, + columns, + output_schema: Arc::new(ArrowSchema::new_with_metadata( + output_fields, + input_schema.metadata().clone(), + )), + }) + } + + pub(crate) fn output_schema(&self) -> &SchemaRef { + &self.output_schema + } + + pub(crate) 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. @@ -1699,46 +2090,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, @@ -7609,6 +7977,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(); 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/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() } 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?; 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);