Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 114 additions & 17 deletions rust/lance-core/src/datatypes/field.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ pub const LANCE_UNENFORCED_CLUSTERING_KEY_POSITION: &str =
/// The value should be non-negative i32 value. Any negative value will be seen as -1.
pub const LANCE_FIELD_ID_KEY: &str = "lance:field_id";

const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"];

fn has_blob_v2_extension(field: &ArrowField) -> bool {
field
.metadata()
Expand Down Expand Up @@ -269,24 +271,15 @@ impl Field {
}

pub fn apply_projection(&self, projection: &Projection) -> Option<Self> {
// Map fields encode their physical layout as a single child entries
// struct (`Struct<key, value>`) whose presence is required for the
// parent to be readable — we never want to filter into that subtree.
// But the parent field itself is still subject to selection: if the
// caller didn't ask for this Map column, drop it like any other
// non-selected leaf. Without this early return the unconditional
// children clone would keep `children.is_empty() == false` forever
// and every Map column in the schema would survive every projection,
// pulling tens-of-bytes-per-row of unrelated data through downstream
// operators (notably `SortExec` in scalar-index training, where it
// was responsible for >100 GiB external-sort spills on real-world
// tables).
if self.logical_type.is_map() && !projection.contains_field_id(self.id) {
// Maps and blob descriptors are atomic physical layouts. Map children
// must remain together, while projected blob descriptor children may
// have synthetic IDs that cannot be selected independently.
let is_atomic_layout = self.logical_type.is_map() || self.is_blob();
if is_atomic_layout && !projection.contains_field_id(self.id) {
return None;
}

let children = if self.logical_type.is_map() {
// Map field is selected: keep all children intact.
let children = if is_atomic_layout {
self.children.clone()
} else {
self.children
Expand Down Expand Up @@ -549,6 +542,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:
Expand All @@ -575,6 +582,31 @@ 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)
.expect("LargeBinary is always a valid logical type");
self.children.clear();
self.encoding = Some(Encoding::VarBinary);
if is_blob_v2 {
self.metadata.remove(BLOB_META_KEY);
for key in PACKED_KEYS {
self.metadata.remove(key);
}
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() {
Expand Down Expand Up @@ -812,6 +844,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
Expand Down Expand Up @@ -847,7 +886,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
Expand Down Expand Up @@ -1038,7 +1077,6 @@ impl Field {

// Check if field has metadata `packed` set to true, this check is case insensitive.
pub fn is_packed_struct(&self) -> bool {
const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"];
PACKED_KEYS.iter().any(|key| {
self.metadata
.get(*key)
Expand Down Expand Up @@ -1847,13 +1885,27 @@ 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()
.unwrap();
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())]);
Expand All @@ -1874,6 +1926,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]
Expand Down Expand Up @@ -1944,4 +2002,43 @@ mod tests {
.unwrap();
assert_eq!(unloaded_projected, unloaded);
}

#[test]
fn blob_descriptor_projection_preserves_synthetic_children() {
let metadata =
HashMap::from([(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string())]);
let mut blob: Field = ArrowField::new(
"blob",
DataType::Struct(
vec![
ArrowField::new("data", DataType::LargeBinary, true),
ArrowField::new("uri", DataType::Utf8, true),
]
.into(),
),
true,
)
.with_metadata(metadata)
.try_into()
.unwrap();
let mut next_id = 0;
blob.set_id(-1, &mut next_id);

let schema = Arc::new(crate::datatypes::Schema {
fields: vec![blob],
metadata: HashMap::new(),
});
let descriptor_schema = Projection::full(schema)
.with_blob_handling(crate::datatypes::BlobHandling::BlobsDescriptions)
.to_bare_schema();
assert!(
descriptor_schema.fields[0]
.children
.iter()
.all(|child| child.id == -1)
);

let projected = Projection::full(Arc::new(descriptor_schema)).to_bare_schema();
assert_eq!(projected.fields[0].children.len(), 5);
}
}
26 changes: 26 additions & 0 deletions rust/lance-core/src/datatypes/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -1095,10 +1106,25 @@ impl BlobHandling {
self.should_unload(field)
}

/// 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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Filtering a blob field can still panic on this path. BlobsDescriptions replaces a blob v2 leaf with descriptor children whose IDs are all -1; the subsequent union/projection drops those children, so applying the projection selects the parent with no children and hits the assertion in Field::apply_projection. For example, blobs IS NOT NULL on the new List<Blob> path, or image.image_bytes IS NOT NULL, raises pyo3_runtime.PanicException. The filtered-read test only filters id, so this case is not covered.

.collect();
field
}
}
Expand Down
9 changes: 2 additions & 7 deletions rust/lance-encoding/src/encodings/logical/blob.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Vec<EncodeTask>> {
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")
Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading