diff --git a/python/python/tests/test_schema.py b/python/python/tests/test_schema.py index fcff283ebe2..c384466082f 100644 --- a/python/python/tests/test_schema.py +++ b/python/python/tests/test_schema.py @@ -6,6 +6,7 @@ import lance import pyarrow as pa +import pytest from lance.schema import LanceSchema @@ -60,3 +61,15 @@ def test_lance_schema(tmp_path: Path): s_fields = fields[1].children() assert s_fields[0].name() == "new_name" assert s_fields[0].id() == 2 + + +def test_lance_schema_from_protos_rejects_missing_parent(): + # name (field 2): child; id (field 3): 7; parent_id (field 4): 42; + # logical_type (field 5): int32. + field_proto = b"\x12\x05child\x18\x07\x20\x2a\x2a\x05int32" + + with pytest.raises( + ValueError, + match="Field 'child' \\(id=7\\) references parent id 42", + ): + LanceSchema._from_protos("{}", field_proto) diff --git a/python/src/schema.rs b/python/src/schema.rs index db4f7369710..8cdc2115cd1 100644 --- a/python/src/schema.rs +++ b/python/src/schema.rs @@ -179,7 +179,8 @@ impl LanceSchema { fields: Fields(fields), metadata, }; - let schema = Schema::from(fields_with_meta); + let schema = Schema::try_from(fields_with_meta) + .map_err(|err| PyValueError::new_err(format!("Failed to reconstruct schema: {err}")))?; Ok(Self(schema)) } diff --git a/rust/lance-file/Cargo.toml b/rust/lance-file/Cargo.toml index f08cd3457aa..56536e840e9 100644 --- a/rust/lance-file/Cargo.toml +++ b/rust/lance-file/Cargo.toml @@ -61,5 +61,9 @@ features = ["protoc"] name = "reader" harness = false +[[bench]] +name = "schema" +harness = false + [lints] workspace = true diff --git a/rust/lance-file/benches/schema.rs b/rust/lance-file/benches/schema.rs new file mode 100644 index 00000000000..b8da23d9777 --- /dev/null +++ b/rust/lance-file/benches/schema.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::hint::black_box; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use lance_core::datatypes::Schema; +use lance_file::{datatypes::Fields, format::pb}; + +fn proto_field(id: i32, parent_id: i32, name: String, logical_type: &str) -> pb::Field { + pb::Field { + id, + parent_id, + name, + logical_type: logical_type.to_owned(), + ..Default::default() + } +} + +/// Builds a pre-order flat schema with `num_physical_columns` physical leaves. +/// +/// Each struct contributes one parent and two `int32` leaves. Root fields use +/// `-1` as `parent_id`, and each struct consumes a three-ID block. +fn wide_two_leaf_structs(num_physical_columns: usize) -> Fields { + assert_eq!(num_physical_columns % 2, 0); + let num_structs = num_physical_columns / 2; + let mut fields = Vec::with_capacity(num_structs + num_physical_columns); + + for struct_index in 0..num_structs { + let parent_id = (struct_index * 3) as i32; + fields.push(proto_field( + parent_id, + -1, + format!("struct_{struct_index}"), + "struct", + )); + fields.push(proto_field( + parent_id + 1, + parent_id, + format!("left_{struct_index}"), + "int32", + )); + fields.push(proto_field( + parent_id + 2, + parent_id, + format!("right_{struct_index}"), + "int32", + )); + } + + Fields(fields) +} + +fn bench_schema_reconstruction(c: &mut Criterion) { + let mut group = c.benchmark_group("schema_from_flat_fields"); + + for num_physical_columns in [1024, 4096, 16_384, 65_536] { + let fields = wide_two_leaf_structs(num_physical_columns); + group.throughput(Throughput::Elements(fields.0.len() as u64)); + group.bench_with_input( + BenchmarkId::new("physical_columns", num_physical_columns), + &fields, + |bencher, fields| { + bencher.iter(|| Schema::try_from(black_box(fields)).unwrap()); + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_schema_reconstruction); +criterion_main!(benches); diff --git a/rust/lance-file/src/datatypes.rs b/rust/lance-file/src/datatypes.rs index ac6a8d7b293..c31cb5c97e7 100644 --- a/rust/lance-file/src/datatypes.rs +++ b/rust/lance-file/src/datatypes.rs @@ -99,6 +99,32 @@ impl From<&Field> for pb::Field { pub struct Fields(pub Vec); +struct FieldNode { + field: Field, + child_indices: Vec, +} + +/// Searches in pre-order depth-first order and returns the first matching node, +/// preserving the legacy parent tie-break for duplicate field IDs. +fn first_field_index_by_id( + nodes: &[FieldNode], + root_indices: &[usize], + field_id: i32, +) -> Option { + let mut to_visit = Vec::with_capacity(nodes.len()); + to_visit.extend(root_indices.iter().rev().copied()); + + while let Some(node_index) = to_visit.pop() { + let node = &nodes[node_index]; + if node.field.id == field_id { + return Some(node_index); + } + to_visit.extend(node.child_indices.iter().rev().copied()); + } + + None +} + impl From<&Field> for Fields { fn from(field: &Field) -> Self { let mut protos = vec![pb::Field::from(field)]; @@ -107,24 +133,119 @@ impl From<&Field> for Fields { } } -/// Convert list of protobuf `Field` to a Schema. -impl From<&Fields> for Schema { - fn from(fields: &Fields) -> Self { - let mut schema = Self { - fields: vec![], - metadata: HashMap::default(), - }; - - fields.0.iter().for_each(|f| { - if f.parent_id == -1 { - schema.fields.push(Field::from(f)); +/// Reconstruct a schema from a flat, pre-order protobuf field list. +/// +/// Parent fields must appear before their children. Historical manifests may +/// contain duplicate field IDs, so an ID may not identify a unique parent. For +/// those references, reconstruction preserves the legacy +/// [`Schema::mut_field_by_id`] tie-break by selecting the first matching field +/// in pre-order depth-first traversal. +/// +/// # Examples +/// +/// ``` +/// use lance_core::datatypes::Schema; +/// use lance_file::{datatypes::Fields, format::pb}; +/// +/// let field = pb::Field { +/// id: 0, +/// parent_id: -1, +/// name: "value".to_owned(), +/// logical_type: "int32".to_owned(), +/// ..Default::default() +/// }; +/// let fields = Fields(vec![field]); +/// let schema = Schema::try_from(&fields)?; +/// assert_eq!(schema.fields[0].name, "value"); +/// # Ok::<(), lance_core::Error>(()) +/// ``` +impl TryFrom<&Fields> for Schema { + type Error = Error; + + fn try_from(fields: &Fields) -> Result { + let mut nodes: Vec = Vec::with_capacity(fields.0.len()); + let mut root_indices = Vec::with_capacity(fields.0.len()); + let mut field_indices: HashMap> = HashMap::with_capacity(fields.0.len()); + + for proto_field in &fields.0 { + let parent_index = if proto_field.parent_id == -1 { + None } else { - let parent = schema.mut_field_by_id(f.parent_id).unwrap(); - parent.children.push(Field::from(f)); + let parent_index = match field_indices.get(&proto_field.parent_id) { + Some(Some(parent_index)) => *parent_index, + Some(None) => { + // Duplicate IDs are invalid but occur in historical + // manifests. Match the legacy tree traversal only for + // these ambiguous parent references so valid schemas + // retain the linear fast path. + first_field_index_by_id(&nodes, &root_indices, proto_field.parent_id) + .ok_or_else(|| { + Error::internal(format!( + "Duplicate field id {} has no existing arena node", + proto_field.parent_id + )) + })? + } + None => { + return Err(Error::schema(format!( + "Field '{}' (id={}) references parent id {}, which must appear earlier in the protobuf field list", + proto_field.name, proto_field.id, proto_field.parent_id + ))); + } + }; + Some(parent_index) + }; + + let node_index = nodes.len(); + if let Some(parent_index) = parent_index { + nodes[parent_index].child_indices.push(node_index); + } else { + root_indices.push(node_index); } - }); + nodes.push(FieldNode { + field: Field::from(proto_field), + child_indices: Vec::new(), + }); + + field_indices + .entry(proto_field.id) + .and_modify(|field_index| *field_index = None) + .or_insert(Some(node_index)); + } - schema + let mut fields_by_node = Vec::with_capacity(nodes.len()); + fields_by_node.resize_with(nodes.len(), || None); + for (node_index, mut node) in nodes.into_iter().enumerate().rev() { + node.field.children.reserve(node.child_indices.len()); + for child_index in node.child_indices { + let child = fields_by_node + .get_mut(child_index) + .and_then(Option::take) + .ok_or_else(|| { + Error::internal(format!( + "Schema field arena node {child_index} was not materialized before its parent" + )) + })?; + node.field.children.push(child); + } + fields_by_node[node_index] = Some(node.field); + } + + let fields = root_indices + .into_iter() + .map(|root_index| { + fields_by_node[root_index].take().ok_or_else(|| { + Error::internal(format!( + "Schema field arena root node {root_index} was not materialized" + )) + }) + }) + .collect::>>()?; + + Ok(Self { + fields, + metadata: HashMap::default(), + }) } } @@ -133,9 +254,28 @@ pub struct FieldsWithMeta { pub metadata: HashMap>, } -/// Convert list of protobuf `Field` and Metadata to a Schema. -impl From for Schema { - fn from(fields_with_meta: FieldsWithMeta) -> Self { +/// Reconstruct a schema from flat protobuf fields and schema metadata. +/// +/// # Examples +/// +/// ``` +/// use std::collections::HashMap; +/// +/// use lance_core::datatypes::Schema; +/// use lance_file::datatypes::{Fields, FieldsWithMeta}; +/// +/// let fields = FieldsWithMeta { +/// fields: Fields(Vec::new()), +/// metadata: HashMap::from([("owner".to_owned(), b"lance".to_vec())]), +/// }; +/// let schema = Schema::try_from(fields)?; +/// assert_eq!(schema.metadata["owner"], "lance"); +/// # Ok::<(), lance_core::Error>(()) +/// ``` +impl TryFrom for Schema { + type Error = Error; + + fn try_from(fields_with_meta: FieldsWithMeta) -> Result { let lance_metadata = fields_with_meta .metadata .into_iter() @@ -145,11 +285,11 @@ impl From for Schema { }) .collect(); - let schema_with_fields = Self::from(&fields_with_meta.fields); - Self { + let schema_with_fields = Self::try_from(&fields_with_meta.fields)?; + Ok(Self { fields: schema_with_fields.fields, metadata: lance_metadata, - } + }) } } @@ -270,14 +410,27 @@ pub async fn populate_schema_dictionary(schema: &mut Schema, reader: &dyn Reader #[cfg(test)] mod tests { + use std::collections::HashMap; + use arrow_schema::DataType; use arrow_schema::Field as ArrowField; use arrow_schema::Fields as ArrowFields; use arrow_schema::Schema as ArrowSchema; + use lance_core::Error; use lance_core::datatypes::Schema; - use std::collections::HashMap; use super::{Fields, FieldsWithMeta}; + use crate::format::pb; + + fn proto_field(id: i32, parent_id: i32, name: String, logical_type: &str) -> pb::Field { + pb::Field { + id, + parent_id, + name, + logical_type: logical_type.to_owned(), + ..Default::default() + } + } #[test] fn test_schema_set_ids() { @@ -317,10 +470,120 @@ mod tests { let expected_schema = Schema::try_from(&arrow_schema).unwrap(); let fields_with_meta: FieldsWithMeta = (&expected_schema).into(); - let schema = Schema::from(fields_with_meta); + let schema = Schema::try_from(fields_with_meta).unwrap(); assert_eq!(expected_schema, schema); } + #[test] + fn test_reconstruct_wide_nested_schema() { + const NUM_STRUCTS: usize = 4096; + + let mut proto_fields = Vec::with_capacity(NUM_STRUCTS * 3); + for struct_index in 0..NUM_STRUCTS { + let parent_id = (struct_index * 3) as i32; + proto_fields.push(proto_field( + parent_id, + -1, + format!("struct_{struct_index}"), + "struct", + )); + proto_fields.push(proto_field( + parent_id + 1, + parent_id, + format!("left_{struct_index}"), + "int32", + )); + proto_fields.push(proto_field( + parent_id + 2, + parent_id, + format!("right_{struct_index}"), + "int32", + )); + } + + let fields = Fields(proto_fields); + let schema = Schema::try_from(&fields).unwrap(); + assert_eq!(schema.fields.len(), NUM_STRUCTS); + for (struct_index, field) in schema.fields.iter().enumerate() { + let parent_id = (struct_index * 3) as i32; + assert_eq!(field.id, parent_id); + assert_eq!(field.name, format!("struct_{struct_index}")); + assert_eq!(field.children.len(), 2); + assert_eq!(field.children[0].id, parent_id + 1); + assert_eq!(field.children[0].name, format!("left_{struct_index}")); + assert_eq!(field.children[1].id, parent_id + 2); + assert_eq!(field.children[1].name, format!("right_{struct_index}")); + } + } + + #[test] + fn test_reconstruct_deep_nested_schema() { + const DEPTH: usize = 1024; + + let proto_fields = (0..DEPTH) + .map(|depth| { + proto_field( + depth as i32, + if depth == 0 { -1 } else { depth as i32 - 1 }, + format!("level_{depth}"), + if depth + 1 == DEPTH { + "int32" + } else { + "struct" + }, + ) + }) + .collect(); + + let fields = Fields(proto_fields); + let schema = Schema::try_from(&fields).unwrap(); + assert_eq!(schema.fields.len(), 1); + let mut field = &schema.fields[0]; + for depth in 0..DEPTH { + assert_eq!(field.id, depth as i32); + assert_eq!(field.name, format!("level_{depth}")); + if depth + 1 == DEPTH { + assert!(field.children.is_empty()); + } else { + assert_eq!(field.children.len(), 1); + field = &field.children[0]; + } + } + } + + #[test] + fn test_reconstruct_schema_reports_missing_parent() { + let fields = Fields(vec![proto_field(7, 42, "child".to_owned(), "int32")]); + + let error = Schema::try_from(&fields).unwrap_err(); + assert!(matches!(&error, Error::Schema { .. })); + assert!( + error.to_string().contains( + "Field 'child' (id=7) references parent id 42, which must appear earlier" + ) + ); + } + + #[test] + fn test_reconstruct_schema_preserves_legacy_duplicate_id_match() { + let fields = Fields(vec![ + proto_field(1, -1, "root_a".to_owned(), "struct"), + proto_field(2, -1, "root_b".to_owned(), "struct"), + proto_field(2, 1, "nested_duplicate".to_owned(), "struct"), + proto_field(3, 2, "child".to_owned(), "int32"), + ]); + + let schema = Schema::try_from(&fields).unwrap(); + assert_eq!(schema.fields.len(), 2); + assert_eq!(schema.fields[0].name, "root_a"); + assert_eq!(schema.fields[0].children.len(), 1); + assert_eq!(schema.fields[0].children[0].name, "nested_duplicate"); + assert_eq!(schema.fields[0].children[0].children.len(), 1); + assert_eq!(schema.fields[0].children[0].children[0].name, "child"); + assert_eq!(schema.fields[1].name, "root_b"); + assert!(schema.fields[1].children.is_empty()); + } + #[test] fn test_clustering_key_roundtrip() { let arrow_schema = ArrowSchema::new(vec![ @@ -351,7 +614,7 @@ mod tests { // Round-trip through protobuf let fields_with_meta: FieldsWithMeta = (&schema).into(); - let restored = Schema::from(fields_with_meta); + let restored = Schema::try_from(fields_with_meta).unwrap(); let ck2 = restored.unenforced_clustering_key(); assert_eq!(ck2.len(), 2); diff --git a/rust/lance-file/src/previous/format/metadata.rs b/rust/lance-file/src/previous/format/metadata.rs index 11ba00c3243..209d0733cec 100644 --- a/rust/lance-file/src/previous/format/metadata.rs +++ b/rust/lance-file/src/previous/format/metadata.rs @@ -62,10 +62,10 @@ impl TryFrom for Metadata { manifest_position: Some(m.manifest_position as usize), stats_metadata: if let Some(stats_meta) = m.statistics { Some(StatisticsMetadata { - schema: Schema::from(FieldsWithMeta { + schema: Schema::try_from(FieldsWithMeta { fields: Fields(stats_meta.schema), metadata: Default::default(), - }), + })?, leaf_field_ids: stats_meta.fields, page_table_position: stats_meta.page_table_position as usize, }) diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index cea89235d73..2edd858b07b 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -961,7 +961,7 @@ impl FileReader { fields: Fields(pb_schema.fields), metadata: pb_schema.metadata, }; - let schema = lance_core::datatypes::Schema::from(fields_with_meta); + let schema = Schema::try_from(fields_with_meta)?; Ok((num_rows, schema)) } diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index cd0a403621f..83c9643e600 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -900,7 +900,7 @@ impl TryFrom for Manifest { Some(format) => DataStorageFormat::from(format), }; - let schema = Schema::from(fields_with_meta); + let schema = Schema::try_from(fields_with_meta)?; Ok(Self { schema, diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 372adee27db..b5ec6973c37 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -3246,7 +3246,7 @@ impl TryFrom for Transaction { .into_iter() .map(Fragment::try_from) .collect::>>()?, - schema: Schema::from(&Fields(schema)), + schema: Schema::try_from(&Fields(schema))?, config_upsert_values: config_upsert_option, initial_bases: if initial_bases.is_empty() { None @@ -3314,7 +3314,7 @@ impl TryFrom for Transaction { .into_iter() .map(Fragment::try_from) .collect::>>()?, - schema: Schema::from(&Fields(schema)), + schema: Schema::try_from(&Fields(schema))?, }, Some(pb::transaction::Operation::Restore(pb::transaction::Restore { version })) => { Operation::Restore { version } @@ -3368,7 +3368,7 @@ impl TryFrom for Transaction { }, Some(pb::transaction::Operation::Project(pb::transaction::Project { schema })) => { Operation::Project { - schema: Schema::from(&Fields(schema)), + schema: Schema::try_from(&Fields(schema))?, } } Some(pb::transaction::Operation::UpdateConfig(update_config)) => {