diff --git a/rust/lance-core/src/datatypes/field.rs b/rust/lance-core/src/datatypes/field.rs index 9f06d421949..c76923958f9 100644 --- a/rust/lance-core/src/datatypes/field.rs +++ b/rust/lance-core/src/datatypes/field.rs @@ -549,6 +549,20 @@ impl Field { .get(ARROW_EXT_NAME_KEY) .map(|name| name == BLOB_V2_EXT_NAME) .unwrap_or(false) + || self.is_blob_v2_descriptor() + } + + fn is_blob_v2_descriptor(&self) -> bool { + self.metadata.contains_key(BLOB_META_KEY) + && self.logical_type == BLOB_V2_DESC_LANCE_FIELD.logical_type + && self.children.len() == BLOB_V2_DESC_LANCE_FIELD.children.len() + && self + .children + .iter() + .zip(BLOB_V2_DESC_LANCE_FIELD.children.iter()) + .all(|(child, expected)| { + child.name == expected.name && child.data_type() == expected.data_type() + }) } // Blob columns intentionally have two schema representations: @@ -575,6 +589,29 @@ impl Field { } } + /// Convert a blob field to the materialized binary payload view. + /// + /// The field keeps its name and id but uses `LargeBinary` with no children. + /// Blob v2 fields retain their extension marker internally so scan planning + /// can recognize the binary view before exposing a plain Arrow binary field. + pub fn binary_blob_mut(&mut self) { + if !self.is_blob() { + return; + } + let is_blob_v2 = self.is_blob_v2(); + + self.logical_type = LogicalType::try_from(&DataType::LargeBinary).unwrap(); + self.children.clear(); + self.encoding = Some(Encoding::VarBinary); + if is_blob_v2 { + self.metadata.remove(BLOB_META_KEY); + self.metadata.remove("packed"); + self.metadata.remove("lance-encoding:packed"); + self.metadata + .insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + } + } + /// Convert blob v2 fields in this field tree to their descriptor view. pub fn unload_blobs_recursive(&mut self) { if self.is_blob_v2() { @@ -812,6 +849,13 @@ impl Field { } if self.is_blob() != other.is_blob() { + if ignore_types { + return Ok(if self.id >= 0 { + self.clone() + } else { + other.clone() + }); + } return Err(Error::arrow(format!( "Attempt to intersect blob and non-blob field: {}", self.name @@ -847,7 +891,7 @@ impl Field { .iter() .filter_map(|c| { if let Some(other_child) = other.child(&c.name) { - let intersection = c.intersection(other_child).ok()?; + let intersection = c.do_intersection(other_child, ignore_types).ok()?; Some(intersection) } else { None @@ -1847,6 +1891,14 @@ mod tests { #[test] fn blob_unloaded_mut_selects_layout_from_metadata() { let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]); + let mut binary_field: Field = ArrowField::new("blob", DataType::LargeBinary, true) + .with_metadata(metadata.clone()) + .try_into() + .unwrap(); + binary_field.binary_blob_mut(); + assert!(binary_field.metadata.contains_key(BLOB_META_KEY)); + assert!(!binary_field.is_blob_v2()); + let mut field: Field = ArrowField::new("blob", DataType::LargeBinary, true) .with_metadata(metadata) .try_into() @@ -1854,6 +1906,12 @@ mod tests { field.unloaded_mut(); assert_eq!(field.children.len(), 2); assert_eq!(field.logical_type, BLOB_DESC_LANCE_FIELD.logical_type); + assert!(field.is_blob()); + assert!(!field.is_blob_v2()); + field.unloaded_mut(); + assert_eq!(field.children.len(), 2); + assert_eq!(field.logical_type, BLOB_DESC_LANCE_FIELD.logical_type); + assert!(!field.is_blob_v2()); let metadata = HashMap::from([(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string())]); @@ -1874,6 +1932,12 @@ mod tests { field.unloaded_mut(); assert_eq!(field.children.len(), 5); assert_eq!(field.logical_type, BLOB_V2_DESC_LANCE_FIELD.logical_type); + assert!(!field.metadata.contains_key(ARROW_EXT_NAME_KEY)); + assert!(field.is_blob_v2()); + field.unloaded_mut(); + assert_eq!(field.children.len(), 5); + assert_eq!(field.logical_type, BLOB_V2_DESC_LANCE_FIELD.logical_type); + assert!(!field.metadata.contains_key(ARROW_EXT_NAME_KEY)); } #[test] diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 7f2cbc02f07..9d4ee3ddf87 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -1071,6 +1071,17 @@ pub enum BlobHandling { } impl BlobHandling { + fn should_load_binary(&self, field: &Field) -> bool { + if !field.is_blob() { + return false; + } + match self { + Self::AllBinary => true, + Self::SomeBlobsBinary(set) | Self::SomeBinary(set) => set.contains(&(field.id as u32)), + Self::BlobsDescriptions | Self::AllDescriptions => false, + } + } + fn should_unload(&self, field: &Field) -> bool { // Blob v2 columns are Structs, so we need to treat any blob-marked field as unloadable // even if the physical data type is not binary-like. @@ -1086,10 +1097,25 @@ impl BlobHandling { } } + /// Apply this blob handling policy to a projected field tree. + /// + /// Blob descriptor modes convert blob leaves to descriptor views. Binary + /// modes convert selected blob leaves to `LargeBinary`. Non-blob nested + /// fields are preserved while their children are handled recursively. pub fn unload_if_needed(&self, mut field: Field) -> Field { + if self.should_load_binary(&field) { + field.binary_blob_mut(); + return field; + } if self.should_unload(&field) { field.unloaded_mut(); + return field; } + field.children = field + .children + .into_iter() + .map(|child| self.unload_if_needed(child)) + .collect(); field } } diff --git a/rust/lance-encoding/src/encodings/logical/blob.rs b/rust/lance-encoding/src/encodings/logical/blob.rs index cad2112bafe..d0b90cd0a22 100644 --- a/rust/lance-encoding/src/encodings/logical/blob.rs +++ b/rust/lance-encoding/src/encodings/logical/blob.rs @@ -267,16 +267,11 @@ impl FieldEncoder for BlobV2StructuralEncoder { &mut self, array: ArrayRef, external_buffers: &mut OutOfLineBuffers, - mut repdef: RepDefBuilder, + repdef: RepDefBuilder, row_number: u64, num_rows: u64, ) -> Result> { let struct_arr = array.as_struct(); - if let Some(validity) = struct_arr.nulls() { - repdef.add_validity_bitmap(validity.clone()); - } else { - repdef.add_no_null(struct_arr.len()); - } let kind_col = struct_arr .column_by_name("kind") @@ -403,7 +398,7 @@ impl FieldEncoder for BlobV2StructuralEncoder { let descriptor_array = Arc::new(StructArray::try_new( BLOB_V2_DESC_FIELDS.clone(), children, - None, + struct_arr.nulls().cloned(), )?) as ArrayRef; self.descriptor_encoder.maybe_encode( diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 84c1b8bdcad..6bc5ea8f6ab 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -12,16 +12,19 @@ use std::{ use arrow::array::AsArray; use arrow::datatypes::{UInt8Type, UInt32Type, UInt64Type}; -use arrow_array::RecordBatch; -use arrow_array::{Array, ArrayRef}; -use arrow_schema::{DataType as ArrowDataType, Field as ArrowField}; +use arrow_array::{ + Array, ArrayRef, GenericListArray, OffsetSizeTrait, RecordBatch, builder::LargeBinaryBuilder, +}; +use arrow_buffer::{OffsetBuffer, ScalarBuffer}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use bytes::Bytes; use futures::future::BoxFuture; use futures::stream::BoxStream; use futures::{FutureExt, StreamExt, TryStreamExt, stream}; use lance_arrow::{ - BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY, - BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, FieldExt, r#struct::StructArrayExt, + ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY, FieldExt, + list::ListArrayExt, r#struct::StructArrayExt, }; use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry}; use lance_io::scheduler::{FileScheduler, ScanScheduler, SchedulerConfig}; @@ -37,9 +40,9 @@ use crate::blob::{ is_logical_blob_v2_field, is_prepared_blob_v2_field, validate_prepared_blob_array, }; use arrow_array::StructArray; -use lance_core::datatypes::{BlobKind, BlobVersion, parse_field_path}; +use lance_core::datatypes::{BlobKind, BlobVersion, Field as LanceField, Schema, parse_field_path}; use lance_core::utils::blob::blob_path; -use lance_core::{Error, Result, utils::address::RowAddress}; +use lance_core::{Error, ROW_ADDR, ROW_ADDR_FIELD, Result, utils::address::RowAddress}; use lance_io::traits::Reader; use lance_io::utils::CachedFileSize; @@ -323,6 +326,9 @@ enum BlobPreprocessFieldKind { Struct { children: Vec, }, + List { + child: Box, + }, Passthrough, } @@ -374,6 +380,17 @@ impl BlobPreprocessField { } } + if let ArrowDataType::List(child) | ArrowDataType::LargeList(child) = field.data_type() { + let child = Self::new(child.as_ref())?; + if child.requires_preprocessing() { + return Ok(Self { + kind: BlobPreprocessFieldKind::List { + child: Box::new(child), + }, + }); + } + } + Ok(Self { kind: BlobPreprocessFieldKind::Passthrough, }) @@ -659,6 +676,20 @@ impl BlobPreprocessor { self.preprocess_struct_array(array, field.as_ref(), children) .await } + BlobPreprocessFieldKind::List { child } => match field.data_type() { + ArrowDataType::List(_) => { + self.preprocess_list_array::(array, field.as_ref(), child) + .await + } + ArrowDataType::LargeList(_) => { + self.preprocess_list_array::(array, field.as_ref(), child) + .await + } + _ => Err(Error::internal(format!( + "Blob list preprocessor received non-list field '{}'", + field.name() + ))), + }, } } .boxed() @@ -716,6 +747,79 @@ impl BlobPreprocessor { Ok((Arc::new(struct_array), field)) } + async fn preprocess_list_array( + &mut self, + array: ArrayRef, + field: &ArrowField, + child: &BlobPreprocessField, + ) -> Result<(ArrayRef, Arc)> { + let list_arr = array.as_list::(); + let list_arr = if list_arr.null_count() > 0 { + list_arr.filter_garbage_nulls() + } else { + list_arr.clone() + }; + + let first_offset = *list_arr + .offsets() + .first() + .ok_or_else(|| Error::invalid_input("List offsets cannot be empty"))?; + let last_offset = *list_arr + .offsets() + .last() + .ok_or_else(|| Error::invalid_input("List offsets cannot be empty"))?; + let values_len = list_arr.values().len(); + let needs_trim = first_offset != O::zero() + || last_offset.to_usize().ok_or_else(|| { + Error::invalid_input(format!( + "List field '{}' offset does not fit into usize", + field.name() + )) + })? != values_len; + + let (offsets, values) = if needs_trim { + let values = list_arr.trimmed_values(); + let offsets = list_arr + .offsets() + .iter() + .map(|offset| *offset - first_offset) + .collect::>(); + (OffsetBuffer::new(ScalarBuffer::from(offsets)), values) + } else { + (list_arr.offsets().clone(), list_arr.values().clone()) + }; + + let child_field = match field.data_type() { + ArrowDataType::List(child_field) | ArrowDataType::LargeList(child_field) => { + child_field.clone() + } + other => { + return Err(Error::invalid_input(format!( + "Blob list preprocessor expected list field '{}', got {other}", + field.name() + ))); + } + }; + let (new_values, new_child_field) = + self.preprocess_field(child, values, &child_field).await?; + + let list_array = GenericListArray::::try_new( + new_child_field, + offsets, + new_values, + list_arr.nulls().cloned(), + )?; + let field = Arc::new( + ArrowField::new( + field.name(), + list_array.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ); + Ok((Arc::new(list_array), field)) + } + async fn preprocess_blob_array( &mut self, array: ArrayRef, @@ -1935,16 +2039,57 @@ async fn collect_blob_entries_for_selection( /// descriptor `StructArray`, descending through nested struct children for /// dotted paths. fn leaf_descriptor_struct<'a>(batch: &'a RecordBatch, column: &str) -> Result<&'a StructArray> { + let current = leaf_descriptor_array(batch, column)?; + current + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Blob column '{}' expected descriptor struct but got {}", + column, + current.data_type() + ) + .into(), + ) + }) +} + +fn leaf_descriptor_array<'a>(batch: &'a RecordBatch, column: &str) -> Result<&'a dyn Array> { let path = parse_field_path(column)?; - let mut current = batch + let mut current: &dyn Array = batch .column_by_name(&path[0]) - .expect("validate_blob_column ensured column exists") - .as_struct(); + .ok_or_else(|| { + Error::invalid_input(format!( + "Blob column '{}' was not found in descriptor batch", + column + )) + })? + .as_ref(); for segment in &path[1..] { - current = current + let struct_array = current + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "Blob column path '{}' expected struct before segment '{}' but got {}", + column, + segment, + current.data_type() + ) + .into(), + ) + })?; + current = struct_array .column_by_name(segment) - .expect("validate_blob_column ensured all path segments exist") - .as_struct(); + .ok_or_else(|| { + Error::invalid_input(format!( + "Blob column path '{}' missing segment '{}'", + column, segment + )) + })? + .as_ref(); } Ok(current) } @@ -1970,6 +2115,38 @@ fn blob_version_from_descriptions(descriptions: &StructArray) -> Result { + descriptions: &'a StructArray, + kinds: &'a arrow::array::PrimitiveArray, + positions: &'a arrow::array::PrimitiveArray, + sizes: &'a arrow::array::PrimitiveArray, + blob_ids: &'a arrow::array::PrimitiveArray, + blob_uris: &'a arrow::array::GenericStringArray, +} + +impl<'a> BlobV2DescriptorColumns<'a> { + fn new(descriptions: &'a StructArray) -> Self { + Self { + descriptions, + kinds: descriptions.column(0).as_primitive::(), + positions: descriptions.column(1).as_primitive::(), + sizes: descriptions.column(2).as_primitive::(), + blob_ids: descriptions.column(3).as_primitive::(), + blob_uris: descriptions.column(4).as_string::(), + } + } + + fn is_null_blob(&self, idx: usize) -> Result { + if self.descriptions.is_null(idx) || self.kinds.is_null(idx) { + return Ok(true); + } + let kind = BlobKind::try_from(self.kinds.value(idx))?; + Ok(matches!(kind, BlobKind::Inline) + && self.positions.value(idx) == 0 + && self.sizes.value(idx) == 0) + } +} + /// Convert blob v1 descriptors into logical blob entries. fn collect_blob_entries_v1( dataset: &Arc, @@ -2032,147 +2209,537 @@ async fn collect_blob_entries_v2( descriptions: &StructArray, row_addrs: &arrow::array::PrimitiveArray, ) -> Result> { - let kinds = descriptions.column(0).as_primitive::(); - let positions = descriptions.column(1).as_primitive::(); - let sizes = descriptions.column(2).as_primitive::(); - let blob_ids = descriptions.column(3).as_primitive::(); - let blob_uris = descriptions.column(4).as_string::(); - + let columns = BlobV2DescriptorColumns::new(descriptions); let mut files = Vec::with_capacity(row_addrs.len()); - let mut fragment_cache = HashMap::::new(); - let mut store_cache = HashMap::>::new(); - let mut external_base_path_cache = HashMap::::new(); - let mut source_cache = HashMap::>::new(); + let mut read_context = BlobV2ReadContext::new(dataset, blob_field_id); for (selection_index, row_addr) in row_addrs.values().iter().enumerate() { - let idx = selection_index; - let kind = BlobKind::try_from(kinds.value(idx))?; + if let Some(entry) = read_context + .collect_entry(&columns, selection_index, selection_index, *row_addr) + .await? + { + files.push(entry); + } + } + + Ok(files) +} + +fn is_blob_v2_binary_view(field: &LanceField) -> bool { + field.is_blob_v2() && matches!(field.data_type(), ArrowDataType::LargeBinary) +} + +fn public_blob_v2_binary_output_field(mut field: LanceField) -> LanceField { + if is_blob_v2_binary_view(&field) { + field.metadata.remove(ARROW_EXT_NAME_KEY); + } + field.children = field + .children + .into_iter() + .map(public_blob_v2_binary_output_field) + .collect(); + field +} + +/// Return the public Arrow-facing schema for a blob v2 binary scan. +/// +/// Scan planning uses a blob v2 extension marker on `LargeBinary` leaves to +/// identify payloads that need descriptor-based materialization. This helper +/// removes that internal marker before the schema is exposed to callers. +pub fn public_blob_v2_binary_output_schema(schema: &Schema) -> Schema { + Schema { + fields: schema + .fields + .iter() + .cloned() + .map(public_blob_v2_binary_output_field) + .collect(), + metadata: schema.metadata.clone(), + } +} + +fn field_has_blob_v2_binary_view(field: &LanceField) -> bool { + is_blob_v2_binary_view(field) || field.children.iter().any(field_has_blob_v2_binary_view) +} + +/// Return true if the schema contains a blob v2 leaf in binary payload view. +/// +/// This detects the internal `LargeBinary` view created by +/// [`BlobHandling::AllBinary`](lance_core::datatypes::BlobHandling::AllBinary) +/// or selective binary blob handling. +pub fn schema_has_blob_v2_binary_view(schema: &Schema) -> bool { + schema.fields.iter().any(field_has_blob_v2_binary_view) +} - // Struct is non-nullable; null rows are encoded as inline with zero position/size and empty uri - if matches!(kind, BlobKind::Inline) && positions.value(idx) == 0 && sizes.value(idx) == 0 { +fn blob_v2_descriptor_field(mut field: LanceField) -> LanceField { + if is_blob_v2_binary_view(&field) { + field.unloaded_mut(); + return field; + } + + field.children = field + .children + .into_iter() + .map(blob_v2_descriptor_field) + .collect(); + field +} + +/// Convert blob v2 binary-view leaves back to descriptor-view leaves. +/// +/// Readers use this schema to fetch stored blob descriptors first. The scan +/// layer then materializes those descriptors into the caller's binary payload +/// view after row addresses are available. +pub fn blob_v2_descriptor_schema(schema: &Schema) -> Schema { + Schema { + fields: schema + .fields + .iter() + .cloned() + .map(blob_v2_descriptor_field) + .collect(), + metadata: schema.metadata.clone(), + } +} + +/// Materialize blob v2 descriptor arrays in a decoded batch into binary arrays. +/// +/// The input batch must include `_rowaddr`, which is used to resolve packed, +/// dedicated, inline, and external blob payload locations. The returned batch +/// matches `output_schema` with blob v2 binary leaves exposed as plain +/// `LargeBinary` fields. +pub async fn materialize_blob_v2_binary_batch( + dataset: &Arc, + output_schema: &Schema, + batch: RecordBatch, + keep_row_addr: bool, +) -> Result { + let row_addr_idx = batch + .schema() + .column_with_name(ROW_ADDR) + .ok_or_else(|| { + Error::internal(format!( + "_rowaddr column missing from blob v2 binary scan batch, columns: {:?}", + batch + .schema() + .fields() + .iter() + .map(|field| field.name()) + .collect::>() + )) + })? + .0; + let row_addrs = batch + .column(row_addr_idx) + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + let row_addrs: Arc<[u64]> = row_addrs.into(); + + let mut columns = Vec::with_capacity(output_schema.fields.len() + usize::from(keep_row_addr)); + let mut fields = Vec::with_capacity(output_schema.fields.len() + usize::from(keep_row_addr)); + + for field in &output_schema.fields { + let input = batch + .column_by_name(&field.name) + .ok_or_else(|| { + Error::internal(format!( + "blob v2 binary scan batch missing projected column '{}'", + field.name + )) + })? + .clone(); + let materialized = + materialize_blob_v2_binary_array(dataset, field, input, row_addrs.clone()).await?; + columns.push(materialized); + let output_field = public_blob_v2_binary_output_field(field.clone()); + fields.push(ArrowField::from(&output_field)); + } + + if keep_row_addr { + columns.push(batch.column(row_addr_idx).clone()); + fields.push((*ROW_ADDR_FIELD).clone()); + } + + Ok(RecordBatch::try_new( + Arc::new(ArrowSchema::new_with_metadata( + fields, + batch.schema().metadata().clone(), + )), + columns, + )?) +} + +fn materialize_blob_v2_binary_array<'a>( + dataset: &'a Arc, + field: &'a LanceField, + array: ArrayRef, + row_addrs: Arc<[u64]>, +) -> BoxFuture<'a, Result> { + async move { + if is_blob_v2_binary_view(field) { + let descriptions = array.as_struct(); + return materialize_blob_v2_descriptors( + dataset, + field.id as u32, + descriptions, + row_addrs.as_ref(), + ) + .await; + } + + match field.data_type() { + ArrowDataType::Struct(_) => { + let struct_array = array.as_struct(); + let mut children = Vec::with_capacity(field.children.len()); + for (child_field, child_array) in + field.children.iter().zip(struct_array.columns().iter()) + { + children.push( + materialize_blob_v2_binary_array( + dataset, + child_field, + child_array.clone(), + row_addrs.clone(), + ) + .await?, + ); + } + let public_field = public_blob_v2_binary_output_field(field.clone()); + let ArrowDataType::Struct(fields) = public_field.data_type() else { + unreachable!("public output field preserved struct type") + }; + Ok(Arc::new(StructArray::try_new( + fields, + children, + struct_array.nulls().cloned(), + )?) as ArrayRef) + } + ArrowDataType::List(_) => { + let list_array = array.as_list::(); + materialize_blob_v2_list_array::(dataset, field, list_array, row_addrs).await + } + ArrowDataType::LargeList(_) => { + let list_array = array.as_list::(); + materialize_blob_v2_list_array::(dataset, field, list_array, row_addrs).await + } + _ => Ok(array), + } + } + .boxed() +} + +async fn materialize_blob_v2_list_array( + dataset: &Arc, + field: &LanceField, + list_array: &GenericListArray, + row_addrs: Arc<[u64]>, +) -> Result { + 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 while materializing blob v2 binary scan", + 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 while materializing blob v2 binary scan", + 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() + )) + })?; + for _ in start..end { + child_row_addrs.push(row_addr); + } + normalized_offsets.push(O::usize_as(end - values_start)); + } + let child_row_addrs: Arc<[u64]> = child_row_addrs.into(); + let child = field.children.first().ok_or_else(|| { + Error::internal(format!( + "List field '{}' missing child while materializing blob v2 binary scan", + field.name + )) + })?; + let values = list_array.values().slice(values_start, values_len); + let values = materialize_blob_v2_binary_array(dataset, child, values, child_row_addrs).await?; + let child_field = public_blob_v2_binary_output_field(child.clone()); + let list_array = GenericListArray::::try_new( + Arc::new(ArrowField::from(&child_field)), + OffsetBuffer::new(ScalarBuffer::from(normalized_offsets)), + values, + list_array.nulls().cloned(), + )?; + Ok(Arc::new(list_array)) +} + +async fn materialize_blob_v2_descriptors( + dataset: &Arc, + blob_field_id: u32, + descriptions: &StructArray, + row_addrs: &[u64], +) -> Result { + if descriptions.len() != row_addrs.len() { + return Err(Error::internal(format!( + "blob v2 descriptor count {} did not match row address count {}", + descriptions.len(), + row_addrs.len() + ))); + } + match blob_version_from_descriptions(descriptions)? { + BlobVersion::V1 => { + return Err(Error::not_supported( + "Blob v2 binary materialization received a legacy blob descriptor".to_string(), + )); + } + BlobVersion::V2 => {} + } + + let columns = BlobV2DescriptorColumns::new(descriptions); + let mut read_context = BlobV2ReadContext::new(dataset, blob_field_id); + let mut builder = LargeBinaryBuilder::with_capacity(descriptions.len(), 0); + + for (idx, row_addr) in row_addrs.iter().copied().enumerate() { + if descriptions.is_null(idx) || columns.kinds.is_null(idx) { + builder.append_null(); continue; } - match kind { + let kind = BlobKind::try_from(columns.kinds.value(idx))?; + if matches!(kind, BlobKind::Inline) + && columns.positions.value(idx) == 0 + && columns.sizes.value(idx) == 0 + { + builder.append_value([]); + continue; + } + + let entry = read_context + .collect_entry(&columns, idx, idx, row_addr) + .await?; + if let Some(entry) = entry { + let data = entry.file.read().await?; + builder.append_value(data.as_ref()); + } else { + builder.append_null(); + } + } + + Ok(Arc::new(builder.finish())) +} + +struct BlobV2ReadContext<'a> { + dataset: &'a Arc, + blob_field_id: u32, + fragment_cache: HashMap, + store_cache: HashMap>, + external_base_path_cache: HashMap, + source_cache: HashMap>, +} + +impl<'a> BlobV2ReadContext<'a> { + fn new(dataset: &'a Arc, blob_field_id: u32) -> Self { + Self { + dataset, + blob_field_id, + fragment_cache: HashMap::new(), + store_cache: HashMap::new(), + external_base_path_cache: HashMap::new(), + source_cache: HashMap::new(), + } + } + + async fn collect_entry( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result> { + if columns.is_null_blob(idx)? { + return Ok(None); + } + + let kind = BlobKind::try_from(columns.kinds.value(idx))?; + let entry = match kind { BlobKind::Inline => { - let position = positions.value(idx); - let size = sizes.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let source = shared_blob_source( - &mut source_cache, - location.object_store, - &location.data_file_path, - ); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, position, size, BlobKind::Inline, None), - }); + self.collect_inline(columns, idx, selection_index, row_addr) + .await? } BlobKind::Dedicated => { - let blob_id = blob_ids.value(idx); - let size = sizes.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); - let source = shared_blob_source(&mut source_cache, location.object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, 0, size, BlobKind::Dedicated, None), - }); + self.collect_dedicated(columns, idx, selection_index, row_addr) + .await? } BlobKind::Packed => { - let blob_id = blob_ids.value(idx); - let size = sizes.value(idx); - let position = positions.value(idx); - let location = resolve_blob_read_location( - dataset, - blob_field_id, - *row_addr, - &mut fragment_cache, - &mut store_cache, - ) - .await?; - let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); - let source = shared_blob_source(&mut source_cache, location.object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source(source, position, size, BlobKind::Packed, None), - }); + self.collect_packed(columns, idx, selection_index, row_addr) + .await? } BlobKind::External => { - let uri_or_path = blob_uris.value(idx).to_string(); - let position = positions.value(idx); - let size = sizes.value(idx); - let base_id = blob_ids.value(idx); - let (object_store, path) = if base_id == 0 { - let registry = dataset.session.store_registry(); - let params = dataset - .store_params - .as_ref() - .map(|p| Arc::new((**p).clone())) - .unwrap_or_else(|| Arc::new(ObjectStoreParams::default())); - ObjectStore::from_uri_and_params(registry, &uri_or_path, ¶ms).await? - } else { - let object_store = if let Some(store) = store_cache.get(&base_id) { - store.clone() - } else { - let store = dataset.object_store(Some(base_id)).await?; - store_cache.insert(base_id, store.clone()); - store - }; - let base_root = if let Some(path) = external_base_path_cache.get(&base_id) { - path.clone() - } else { - let base = dataset.manifest.base_paths.get(&base_id).ok_or_else(|| { - Error::invalid_input(format!( - "External blob references unknown base_id {}", - base_id - )) - })?; - let path = base.extract_path(dataset.session.store_registry())?; - external_base_path_cache.insert(base_id, path.clone()); - path - }; - let path = join_base_and_relative_path(&base_root, &uri_or_path)?; - (object_store, path) - }; - let size = if size > 0 { - size - } else { - object_store.size(&path).await? - }; - let source = shared_blob_source(&mut source_cache, object_store, &path); - files.push(BlobEntry { - selection_index, - row_address: *row_addr, - file: BlobFile::with_source( - source, - position, - size, - BlobKind::External, - Some(uri_or_path), - ), - }); + self.collect_external(columns, idx, selection_index, row_addr) + .await? } - } + }; + + Ok(Some(entry)) } - Ok(files) + async fn blob_read_location(&mut self, row_addr: u64) -> Result { + resolve_blob_read_location( + self.dataset, + self.blob_field_id, + row_addr, + &mut self.fragment_cache, + &mut self.store_cache, + ) + .await + } + + async fn collect_inline( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let position = columns.positions.value(idx); + let size = columns.sizes.value(idx); + let location = self.blob_read_location(row_addr).await?; + let source = shared_blob_source( + &mut self.source_cache, + location.object_store, + &location.data_file_path, + ); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, position, size, BlobKind::Inline, None), + }) + } + + async fn collect_dedicated( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let blob_id = columns.blob_ids.value(idx); + let size = columns.sizes.value(idx); + let location = self.blob_read_location(row_addr).await?; + let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); + let source = shared_blob_source(&mut self.source_cache, location.object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, 0, size, BlobKind::Dedicated, None), + }) + } + + async fn collect_packed( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let blob_id = columns.blob_ids.value(idx); + let size = columns.sizes.value(idx); + let position = columns.positions.value(idx); + let location = self.blob_read_location(row_addr).await?; + let path = blob_path(&location.data_file_dir, &location.data_file_key, blob_id); + let source = shared_blob_source(&mut self.source_cache, location.object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source(source, position, size, BlobKind::Packed, None), + }) + } + + async fn collect_external( + &mut self, + columns: &BlobV2DescriptorColumns<'_>, + idx: usize, + selection_index: usize, + row_addr: u64, + ) -> Result { + let uri_or_path = columns.blob_uris.value(idx).to_string(); + let position = columns.positions.value(idx); + let size = columns.sizes.value(idx); + let base_id = columns.blob_ids.value(idx); + let (object_store, path) = if base_id == 0 { + let registry = self.dataset.session.store_registry(); + let params = self + .dataset + .store_params + .as_ref() + .map(|p| Arc::new((**p).clone())) + .unwrap_or_else(|| Arc::new(ObjectStoreParams::default())); + ObjectStore::from_uri_and_params(registry, &uri_or_path, ¶ms).await? + } else { + let object_store = if let Some(store) = self.store_cache.get(&base_id) { + store.clone() + } else { + let store = self.dataset.object_store(Some(base_id)).await?; + self.store_cache.insert(base_id, store.clone()); + store + }; + let base_root = if let Some(path) = self.external_base_path_cache.get(&base_id) { + path.clone() + } else { + let base = self + .dataset + .manifest + .base_paths + .get(&base_id) + .ok_or_else(|| { + Error::invalid_input(format!( + "External blob references unknown base_id {}", + base_id + )) + })?; + let path = base.extract_path(self.dataset.session.store_registry())?; + self.external_base_path_cache.insert(base_id, path.clone()); + path + }; + let path = join_base_and_relative_path(&base_root, &uri_or_path)?; + (object_store, path) + }; + let size = if size > 0 { + size + } else { + object_store.size(&path).await? + }; + let source = shared_blob_source(&mut self.source_cache, object_store, &path); + Ok(BlobEntry { + selection_index, + row_address: row_addr, + file: BlobFile::with_source( + source, + position, + size, + BlobKind::External, + Some(uri_or_path), + ), + }) + } } fn normalize_external_absolute_uri(uri: &str) -> Result { @@ -3837,6 +4404,216 @@ mod tests { assert!(null_blobs.is_empty()); } + #[tokio::test] + async fn test_write_and_scan_list_blob_v2_descriptions() { + let test_dir = TempStrDir::default(); + let packed_payload = vec![0x4B; super::INLINE_MAX + 1024]; + + let mut blob_builder = BlobArrayBuilder::new(4); + blob_builder.push_bytes(b"hello").unwrap(); + blob_builder.push_null().unwrap(); + blob_builder.push_bytes(&packed_payload).unwrap(); + blob_builder.push_bytes(b"tail").unwrap(); + let blob_values = blob_builder.finish().unwrap(); + + let item_field = Arc::new(blob_field("item", true)); + let list_array: ArrayRef = Arc::new( + arrow_array::ListArray::try_new( + item_field.clone(), + arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![ + 0i32, 3, 3, 3, 4, + ])), + blob_values, + Some(arrow_buffer::NullBuffer::from(vec![ + true, true, false, true, + ])), + ) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![Field::new( + "blobs", + DataType::List(item_field), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![list_array]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + + let dataset = Arc::new( + Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let descriptions = dataset + .scan() + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let lists = descriptions.column(0).as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 3, 3, 3, 4]); + assert!(lists.is_valid(0)); + assert!(lists.is_valid(1)); + assert!(lists.is_null(2)); + assert!(lists.is_valid(3)); + + let DataType::List(descriptor_field) = lists.data_type() else { + panic!("unexpected list type: {}", lists.data_type()); + }; + assert!(matches!(descriptor_field.data_type(), DataType::Struct(_))); + assert!(!descriptor_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let descriptors = lists.values().as_struct(); + assert_eq!(descriptors.fields().len(), 5); + assert_eq!(descriptors.fields()[0].name(), "kind"); + assert!(descriptors.is_valid(0)); + assert!(descriptors.is_null(1)); + assert!(descriptors.is_valid(2)); + assert!(descriptors.is_valid(3)); + let kinds = descriptors + .column_by_name("kind") + .unwrap() + .as_primitive::(); + assert_eq!(kinds.value(0), BlobKind::Inline as u8); + assert_eq!(kinds.value(2), BlobKind::Packed as u8); + assert_eq!(kinds.value(3), BlobKind::Inline as u8); + + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + let bytes = scanner + .project(&["blobs"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let lists = bytes.column(0).as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 3, 3, 3, 4]); + assert!(lists.is_valid(0)); + assert!(lists.is_valid(1)); + assert!(lists.is_null(2)); + assert!(lists.is_valid(3)); + let DataType::List(value_field) = lists.data_type() else { + panic!("unexpected list type: {}", lists.data_type()); + }; + assert_eq!(value_field.data_type(), &DataType::LargeBinary); + assert!(!value_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let values = lists.values().as_binary::(); + assert_eq!(values.value(0), b"hello"); + assert!(values.is_null(1)); + assert_eq!(values.value(2), packed_payload.as_slice()); + assert_eq!(values.value(3), b"tail"); + } + + #[tokio::test] + async fn test_write_and_scan_struct_nested_list_blob_v2() { + let test_dir = TempStrDir::default(); + + let mut blob_builder = BlobArrayBuilder::new(2); + blob_builder.push_bytes(b"nested").unwrap(); + blob_builder.push_null().unwrap(); + let blob_values = blob_builder.finish().unwrap(); + + let item_field = Arc::new(blob_field("item", true)); + let list_field = Field::new("blobs", DataType::List(item_field.clone()), true); + let list_array: ArrayRef = Arc::new( + arrow_array::ListArray::try_new( + item_field, + arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![0i32, 2, 2])), + blob_values, + None, + ) + .unwrap(), + ); + let info_fields = vec![Field::new("name", DataType::Utf8, false), list_field]; + let info_array: ArrayRef = Arc::new( + StructArray::try_new( + info_fields.clone().into(), + vec![ + Arc::new(StringArray::from(vec!["row-0", "row-1"])) as ArrayRef, + list_array, + ], + None, + ) + .unwrap(), + ); + + let schema = Arc::new(Schema::new(vec![Field::new( + "info", + DataType::Struct(info_fields.into()), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![info_array]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + + let dataset = Arc::new( + Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let descriptions = dataset + .scan() + .project(&["info"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let info = descriptions.column(0).as_struct(); + assert_eq!( + info.column_by_name("name") + .unwrap() + .as_string::() + .value(0), + "row-0" + ); + let lists = info.column_by_name("blobs").unwrap().as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 2, 2]); + let DataType::List(descriptor_field) = lists.data_type() else { + panic!("unexpected nested list type: {}", lists.data_type()); + }; + assert!(matches!(descriptor_field.data_type(), DataType::Struct(_))); + assert!(!descriptor_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let descriptors = lists.values().as_struct(); + assert_eq!(descriptors.fields().len(), 5); + assert!(descriptors.is_valid(0)); + assert!(descriptors.is_null(1)); + + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + let bytes = scanner + .project(&["info"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let info = bytes.column(0).as_struct(); + let lists = info.column_by_name("blobs").unwrap().as_list::(); + assert_eq!(lists.offsets().inner().as_ref(), &[0, 2, 2]); + let DataType::List(value_field) = lists.data_type() else { + panic!("unexpected nested list type: {}", lists.data_type()); + }; + assert_eq!(value_field.data_type(), &DataType::LargeBinary); + assert!(!value_field.metadata().contains_key(ARROW_EXT_NAME_KEY)); + let values = lists.values().as_binary::(); + assert_eq!(values.value(0), b"nested"); + assert!(values.is_null(1)); + } + #[tokio::test] async fn test_nested_blob_v2_requires_v2_2() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 059fa4e2b36..a963059d1cb 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -8,7 +8,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::{ops::Range, sync::Arc}; use arrow_array::RecordBatch; -use arrow_schema::SchemaRef; +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use datafusion::common::runtime::SpawnedTask; use datafusion::common::stats::Precision; use datafusion::error::{DataFusionError, Result as DataFusionResult}; @@ -60,6 +60,13 @@ use crate::dataset::scanner::{ use super::utils::IoMetrics; +fn public_blob_v2_binary_projection_schema(projection: &Projection) -> SchemaRef { + let schema = projection.to_schema(); + let schema = crate::dataset::blob::public_blob_v2_binary_output_schema(&schema); + let schema: ArrowSchema = (&schema).into(); + Arc::new(schema) +} + #[derive(Debug)] pub struct EvaluatedIndex { index_result: IndexExprResult, @@ -385,7 +392,7 @@ impl FilteredReadStream { .try_collect::>() .await?; - let output_schema = Arc::new(options.projection.to_arrow_schema()); + let output_schema = public_blob_v2_binary_projection_schema(&options.projection); let obj_store = dataset.object_store.clone(); // Explicit options take precedence; otherwise fall back to the @@ -417,11 +424,14 @@ impl FilteredReadStream { let fragment_streams = futures::stream::iter(scoped_fragments) .map({ let scan_range_after_filter = scan_range_after_filter.clone(); + let dataset = dataset.clone(); move |scoped_fragment| { let metrics = global_metrics_clone.clone(); let limit = scan_range_after_filter.as_ref().map(|r| r.end); + let dataset = dataset.clone(); SpawnedTask::spawn( - Self::read_fragment(scoped_fragment, metrics, limit).in_current_span(), + Self::read_fragment(dataset, scoped_fragment, metrics, limit) + .in_current_span(), ) .map(|thread_result| thread_result.unwrap()) } @@ -1076,11 +1086,13 @@ impl FilteredReadStream { // Reads a single fragment into a stream of batch tasks #[instrument(name = "read_fragment", level = "debug", skip_all)] async fn read_fragment( + dataset: Arc, mut fragment_read_task: ScopedFragmentRead, global_metrics: Arc, fragment_soft_limit: Option, ) -> Result>> { - let output_schema = Arc::new(fragment_read_task.projection.to_arrow_schema()); + let output_schema = + public_blob_v2_binary_projection_schema(fragment_read_task.projection.as_ref()); if let Some(filter) = &fragment_read_task.filter { let filter_cols = Planner::column_names_in_expr(filter); @@ -1095,10 +1107,22 @@ impl FilteredReadStream { } } - let read_schema = fragment_read_task.projection.to_bare_schema(); + let output_read_schema = Arc::new(fragment_read_task.projection.to_bare_schema()); + let materialize_blob_v2_binary = + crate::dataset::blob::schema_has_blob_v2_binary_view(output_read_schema.as_ref()); + let read_schema = if materialize_blob_v2_binary { + crate::dataset::blob::blob_v2_descriptor_schema(output_read_schema.as_ref()) + } else { + output_read_schema.as_ref().clone() + }; + let mut frag_read_config = fragment_read_task.frag_read_config(); + if materialize_blob_v2_binary { + frag_read_config = frag_read_config.with_row_address(true); + } + let keep_materialized_row_addr = fragment_read_task.projection.with_row_addr; let mut fragment_reader = fragment_read_task .fragment - .open(&read_schema, fragment_read_task.frag_read_config()) + .open(&read_schema, frag_read_config) .await?; if fragment_read_task.with_deleted_rows { @@ -1112,8 +1136,9 @@ impl FilteredReadStream { let physical_filter = fragment_read_task .filter .map(|filter| { - let planner = - Planner::new(Arc::new(fragment_read_task.projection.to_arrow_schema())); + let planner = Planner::new(public_blob_v2_binary_projection_schema( + fragment_read_task.projection.as_ref(), + )); planner.create_physical_expr(&filter) }) .transpose()?; @@ -1136,7 +1161,7 @@ impl FilteredReadStream { let global_metrics = global_metrics.clone(); let fragment_counted = fragment_counted.clone(); let range_tracker = range_tracker.clone(); - batch_fut + let batch_fut = batch_fut .inspect_ok(move |batch| { let num_rows = batch.num_rows(); global_metrics.rows_scanned.add(num_rows); @@ -1151,7 +1176,28 @@ impl FilteredReadStream { global_metrics.ranges_scanned.add(additional_ranges); } }) - .boxed() + .boxed(); + if materialize_blob_v2_binary { + let dataset = dataset.clone(); + let output_read_schema = output_read_schema.clone(); + batch_fut + .and_then(move |batch| { + let dataset = dataset.clone(); + let output_read_schema = output_read_schema.clone(); + async move { + crate::dataset::blob::materialize_blob_v2_binary_batch( + &dataset, + output_read_schema.as_ref(), + batch, + keep_materialized_row_addr, + ) + .await + } + }) + .boxed() + } else { + batch_fut + } }) .zip(futures::stream::repeat(( physical_filter.clone(), @@ -1595,7 +1641,7 @@ impl FilteredReadExec { )); } } - let output_schema = Arc::new(options.projection.to_arrow_schema()); + let output_schema = public_blob_v2_binary_projection_schema(&options.projection); let num_partitions = match options.threading_mode { FilteredReadThreadingMode::OnePartitionMultipleThreads(_) => 1, FilteredReadThreadingMode::MultiplePartitions(n) => n, @@ -1973,7 +2019,7 @@ impl ExecutionPlan for FilteredReadExec { .clone() .union_columns(filter_columns, OnMissing::Error)?; - let read_schema = Arc::new(read_projection.to_arrow_schema()); + let read_schema = public_blob_v2_binary_projection_schema(&read_projection); let planner = Arc::new(Planner::new(read_schema.clone())); let physical_filter = planner.create_physical_expr(filter)?; diff --git a/rust/lance/src/io/exec/scan.rs b/rust/lance/src/io/exec/scan.rs index 3ec63ce04cc..03b24ca76b6 100644 --- a/rust/lance/src/io/exec/scan.rs +++ b/rust/lance/src/io/exec/scan.rs @@ -192,7 +192,17 @@ impl LanceStream { ) -> Result { let scan_metrics = ScanMetrics::new(metrics, partition); let timer = scan_metrics.baseline_metrics.elapsed_compute().timer(); - let project_schema = projection.clone(); + let materialize_blob_v2_binary = + crate::dataset::blob::schema_has_blob_v2_binary_view(projection.as_ref()); + let read_projection = if materialize_blob_v2_binary { + Arc::new(crate::dataset::blob::blob_v2_descriptor_schema( + projection.as_ref(), + )) + } else { + projection.clone() + }; + let project_schema = read_projection; + let output_projection = projection.clone(); let io_parallelism = dataset.object_store.io_parallelism(); // First, use the value specified by the user in the call // Second, use the default from the environment variable, if specified @@ -275,12 +285,15 @@ impl LanceStream { let scan_scheduler_clone = scan_scheduler.clone(); + let materialize_dataset = dataset; + let keep_materialized_row_addr = config.with_row_address; let config_for_stream = config.clone(); let batches = stream::iter(file_fragments.into_iter().enumerate()) .map(move |(priority, file_fragment)| { let project_schema = project_schema.clone(); let scan_scheduler = scan_scheduler.clone(); let config = config_for_stream.clone(); + let force_row_address = materialize_blob_v2_binary; #[allow(clippy::type_complexity)] let frag_task: BoxFuture< Result>>>>, @@ -288,7 +301,7 @@ impl LanceStream { (async move { let mut frag_config = FragReadConfig::default() .with_row_id(config.with_row_id) - .with_row_address(config.with_row_address) + .with_row_address(config.with_row_address || force_row_address) .with_row_last_updated_at_version( config.with_row_last_updated_at_version, ) @@ -349,6 +362,26 @@ impl LanceStream { ) .stream_in_current_span() .boxed(); + let inner_stream = if materialize_blob_v2_binary { + inner_stream + .and_then(move |batch| { + let dataset = materialize_dataset.clone(); + let output_projection = output_projection.clone(); + async move { + crate::dataset::blob::materialize_blob_v2_binary_batch( + &dataset, + output_projection.as_ref(), + batch, + keep_materialized_row_addr, + ) + .await + .map_err(DataFusionError::from) + } + }) + .boxed() + } else { + inner_stream + }; timer.done(); Ok(Self { @@ -482,7 +515,9 @@ impl core::fmt::Debug for LanceStream { impl RecordBatchStream for LanceStream { fn schema(&self) -> SchemaRef { - let mut schema: ArrowSchema = self.projection.as_ref().into(); + let output_projection = + crate::dataset::blob::public_blob_v2_binary_output_schema(self.projection.as_ref()); + let mut schema: ArrowSchema = (&output_projection).into(); if self.config.with_row_id { schema = schema.try_with_column(ROW_ID_FIELD.clone()).unwrap(); } @@ -602,7 +637,9 @@ impl LanceScanExec { projection: Arc, config: LanceScanConfig, ) -> Self { - let mut output_schema: ArrowSchema = projection.as_ref().into(); + let output_projection = + crate::dataset::blob::public_blob_v2_binary_output_schema(projection.as_ref()); + let mut output_schema: ArrowSchema = (&output_projection).into(); if config.with_row_id { output_schema = output_schema.try_with_column(ROW_ID_FIELD.clone()).unwrap();