fix: reconstruct protobuf schema fields in linear time

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)

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

rust/lance-file/Cargo.toml

@@ -61,5 +61,9 @@ features = ["protoc"]
 name = "reader"
 harness = false
 
+[[bench]]
+name = "schema"
+harness = false
+
 [lints]
 workspace = true

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