lance_file/
datatypes.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use arrow_schema::DataType;
5use async_recursion::async_recursion;
6use lance_arrow::ARROW_EXT_NAME_KEY;
7use lance_arrow::DataTypeExt;
8use lance_core::datatypes::{Dictionary, Encoding, Field, LogicalType, Schema};
9use lance_core::{Error, Result};
10use lance_io::traits::Reader;
11use lance_io::utils::{read_binary_array, read_fixed_stride_array};
12use std::collections::HashMap;
13
14use crate::format::pb;
15
16#[allow(clippy::fallible_impl_from)]
17impl From<&pb::Field> for Field {
18    fn from(field: &pb::Field) -> Self {
19        let lance_metadata: HashMap<String, String> = field
20            .metadata
21            .iter()
22            .map(|(key, value)| {
23                let string_value = String::from_utf8_lossy(value).to_string();
24                (key.clone(), string_value)
25            })
26            .collect();
27        let mut lance_metadata = lance_metadata;
28        if !field.extension_name.is_empty() {
29            lance_metadata.insert(ARROW_EXT_NAME_KEY.to_string(), field.extension_name.clone());
30        }
31        Self {
32            name: field.name.clone(),
33            id: field.id,
34            parent_id: field.parent_id,
35            logical_type: LogicalType::from(field.logical_type.as_str()),
36            metadata: lance_metadata,
37            encoding: match field.encoding {
38                1 => Some(Encoding::Plain),
39                2 => Some(Encoding::VarBinary),
40                3 => Some(Encoding::Dictionary),
41                4 => Some(Encoding::RLE),
42                _ => None,
43            },
44            nullable: field.nullable,
45            children: vec![],
46            dictionary: field.dictionary.as_ref().map(Dictionary::from),
47            unenforced_primary_key_position: if field.unenforced_primary_key_position > 0 {
48                Some(field.unenforced_primary_key_position)
49            } else if field.unenforced_primary_key {
50                Some(0)
51            } else {
52                None
53            },
54            unenforced_clustering_key_position: if field.unenforced_clustering_key_position > 0 {
55                Some(field.unenforced_clustering_key_position)
56            } else {
57                None
58            },
59        }
60    }
61}
62
63impl From<&Field> for pb::Field {
64    fn from(field: &Field) -> Self {
65        let pb_metadata = field
66            .metadata
67            .iter()
68            .map(|(key, value)| (key.clone(), value.clone().into_bytes()))
69            .collect();
70        Self {
71            id: field.id,
72            parent_id: field.parent_id,
73            name: field.name.clone(),
74            logical_type: field.logical_type.to_string(),
75            encoding: match field.encoding {
76                Some(Encoding::Plain) => 1,
77                Some(Encoding::VarBinary) => 2,
78                Some(Encoding::Dictionary) => 3,
79                Some(Encoding::RLE) => 4,
80                _ => 0,
81            },
82            nullable: field.nullable,
83            dictionary: field.dictionary.as_ref().map(pb::Dictionary::from),
84            metadata: pb_metadata,
85            extension_name: field
86                .extension_name()
87                .map(|name| name.to_owned())
88                .unwrap_or_default(),
89            r#type: 0,
90            unenforced_primary_key: field.unenforced_primary_key_position.is_some(),
91            unenforced_primary_key_position: field.unenforced_primary_key_position.unwrap_or(0),
92            unenforced_clustering_key: false,
93            unenforced_clustering_key_position: field
94                .unenforced_clustering_key_position
95                .unwrap_or(0),
96        }
97    }
98}
99
100pub struct Fields(pub Vec<pb::Field>);
101
102struct FieldNode {
103    field: Field,
104    child_indices: Vec<usize>,
105}
106
107/// Searches in pre-order depth-first order and returns the first matching node,
108/// preserving the legacy parent tie-break for duplicate field IDs.
109fn first_field_index_by_id(
110    nodes: &[FieldNode],
111    root_indices: &[usize],
112    field_id: i32,
113) -> Option<usize> {
114    let mut to_visit = Vec::with_capacity(nodes.len());
115    to_visit.extend(root_indices.iter().rev().copied());
116
117    while let Some(node_index) = to_visit.pop() {
118        let node = &nodes[node_index];
119        if node.field.id == field_id {
120            return Some(node_index);
121        }
122        to_visit.extend(node.child_indices.iter().rev().copied());
123    }
124
125    None
126}
127
128impl From<&Field> for Fields {
129    fn from(field: &Field) -> Self {
130        let mut protos = vec![pb::Field::from(field)];
131        protos.extend(field.children.iter().flat_map(|val| Self::from(val).0));
132        Self(protos)
133    }
134}
135
136/// Reconstruct a schema from a flat, pre-order protobuf field list.
137///
138/// Parent fields must appear before their children. Historical manifests may
139/// contain duplicate field IDs, so an ID may not identify a unique parent. For
140/// those references, reconstruction preserves the legacy
141/// [`Schema::mut_field_by_id`] tie-break by selecting the first matching field
142/// in pre-order depth-first traversal.
143///
144/// # Examples
145///
146/// ```
147/// use lance_core::datatypes::Schema;
148/// use lance_file::{datatypes::Fields, format::pb};
149///
150/// let field = pb::Field {
151///     id: 0,
152///     parent_id: -1,
153///     name: "value".to_owned(),
154///     logical_type: "int32".to_owned(),
155///     ..Default::default()
156/// };
157/// let fields = Fields(vec![field]);
158/// let schema = Schema::try_from(&fields)?;
159/// assert_eq!(schema.fields[0].name, "value");
160/// # Ok::<(), lance_core::Error>(())
161/// ```
162impl TryFrom<&Fields> for Schema {
163    type Error = Error;
164
165    fn try_from(fields: &Fields) -> Result<Self> {
166        let mut nodes: Vec<FieldNode> = Vec::with_capacity(fields.0.len());
167        let mut root_indices = Vec::with_capacity(fields.0.len());
168        let mut field_indices: HashMap<i32, Option<usize>> = HashMap::with_capacity(fields.0.len());
169
170        for proto_field in &fields.0 {
171            let parent_index = if proto_field.parent_id == -1 {
172                None
173            } else {
174                let parent_index = match field_indices.get(&proto_field.parent_id) {
175                    Some(Some(parent_index)) => *parent_index,
176                    Some(None) => {
177                        // Duplicate IDs are invalid but occur in historical
178                        // manifests. Match the legacy tree traversal only for
179                        // these ambiguous parent references so valid schemas
180                        // retain the linear fast path.
181                        first_field_index_by_id(&nodes, &root_indices, proto_field.parent_id)
182                            .ok_or_else(|| {
183                                Error::internal(format!(
184                                    "Duplicate field id {} has no existing arena node",
185                                    proto_field.parent_id
186                                ))
187                            })?
188                    }
189                    None => {
190                        return Err(Error::schema(format!(
191                            "Field '{}' (id={}) references parent id {}, which must appear earlier in the protobuf field list",
192                            proto_field.name, proto_field.id, proto_field.parent_id
193                        )));
194                    }
195                };
196                Some(parent_index)
197            };
198
199            let node_index = nodes.len();
200            if let Some(parent_index) = parent_index {
201                nodes[parent_index].child_indices.push(node_index);
202            } else {
203                root_indices.push(node_index);
204            }
205            nodes.push(FieldNode {
206                field: Field::from(proto_field),
207                child_indices: Vec::new(),
208            });
209
210            field_indices
211                .entry(proto_field.id)
212                .and_modify(|field_index| *field_index = None)
213                .or_insert(Some(node_index));
214        }
215
216        let mut fields_by_node = Vec::with_capacity(nodes.len());
217        fields_by_node.resize_with(nodes.len(), || None);
218        for (node_index, mut node) in nodes.into_iter().enumerate().rev() {
219            node.field.children.reserve(node.child_indices.len());
220            for child_index in node.child_indices {
221                let child = fields_by_node
222                    .get_mut(child_index)
223                    .and_then(Option::take)
224                    .ok_or_else(|| {
225                        Error::internal(format!(
226                            "Schema field arena node {child_index} was not materialized before its parent"
227                        ))
228                    })?;
229                node.field.children.push(child);
230            }
231            fields_by_node[node_index] = Some(node.field);
232        }
233
234        let fields = root_indices
235            .into_iter()
236            .map(|root_index| {
237                fields_by_node[root_index].take().ok_or_else(|| {
238                    Error::internal(format!(
239                        "Schema field arena root node {root_index} was not materialized"
240                    ))
241                })
242            })
243            .collect::<Result<Vec<_>>>()?;
244
245        Ok(Self {
246            fields,
247            metadata: HashMap::default(),
248        })
249    }
250}
251
252pub struct FieldsWithMeta {
253    pub fields: Fields,
254    pub metadata: HashMap<String, Vec<u8>>,
255}
256
257/// Reconstruct a schema from flat protobuf fields and schema metadata.
258///
259/// # Examples
260///
261/// ```
262/// use std::collections::HashMap;
263///
264/// use lance_core::datatypes::Schema;
265/// use lance_file::datatypes::{Fields, FieldsWithMeta};
266///
267/// let fields = FieldsWithMeta {
268///     fields: Fields(Vec::new()),
269///     metadata: HashMap::from([("owner".to_owned(), b"lance".to_vec())]),
270/// };
271/// let schema = Schema::try_from(fields)?;
272/// assert_eq!(schema.metadata["owner"], "lance");
273/// # Ok::<(), lance_core::Error>(())
274/// ```
275impl TryFrom<FieldsWithMeta> for Schema {
276    type Error = Error;
277
278    fn try_from(fields_with_meta: FieldsWithMeta) -> Result<Self> {
279        let lance_metadata = fields_with_meta
280            .metadata
281            .into_iter()
282            .map(|(key, value)| {
283                let string_value = String::from_utf8_lossy(&value).to_string();
284                (key, string_value)
285            })
286            .collect();
287
288        let schema_with_fields = Self::try_from(&fields_with_meta.fields)?;
289        Ok(Self {
290            fields: schema_with_fields.fields,
291            metadata: lance_metadata,
292        })
293    }
294}
295
296/// Convert a Schema to a list of protobuf Field.
297impl From<&Schema> for Fields {
298    fn from(schema: &Schema) -> Self {
299        let mut protos = vec![];
300        schema.fields.iter().for_each(|f| {
301            protos.extend(Self::from(f).0);
302        });
303        Self(protos)
304    }
305}
306
307/// Convert a Schema to a list of protobuf Field and Metadata
308impl From<&Schema> for FieldsWithMeta {
309    fn from(schema: &Schema) -> Self {
310        let fields = schema.into();
311        let metadata = schema
312            .metadata
313            .clone()
314            .into_iter()
315            .map(|(key, value)| (key, value.into_bytes()))
316            .collect();
317        Self { fields, metadata }
318    }
319}
320
321impl From<&pb::Dictionary> for Dictionary {
322    fn from(proto: &pb::Dictionary) -> Self {
323        Self {
324            offset: proto.offset as usize,
325            length: proto.length as usize,
326            values: None,
327        }
328    }
329}
330
331impl From<&Dictionary> for pb::Dictionary {
332    fn from(d: &Dictionary) -> Self {
333        Self {
334            offset: d.offset as i64,
335            length: d.length as i64,
336        }
337    }
338}
339
340impl From<Encoding> for pb::Encoding {
341    fn from(e: Encoding) -> Self {
342        match e {
343            Encoding::Plain => Self::Plain,
344            Encoding::VarBinary => Self::VarBinary,
345            Encoding::Dictionary => Self::Dictionary,
346            Encoding::RLE => Self::Rle,
347        }
348    }
349}
350
351#[async_recursion]
352async fn load_field_dictionary<'a>(field: &mut Field, reader: &dyn Reader) -> Result<()> {
353    if let DataType::Dictionary(_, value_type) = field.data_type() {
354        assert!(field.dictionary.is_some());
355        if let Some(dict_info) = field.dictionary.as_mut() {
356            use DataType::*;
357            match value_type.as_ref() {
358                _ if value_type.is_binary_like() => {
359                    dict_info.values = Some(
360                        read_binary_array(
361                            reader,
362                            value_type.as_ref(),
363                            true, // Empty values are null
364                            dict_info.offset,
365                            dict_info.length,
366                            ..,
367                        )
368                        .await?,
369                    );
370                }
371                Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 => {
372                    dict_info.values = Some(
373                        read_fixed_stride_array(
374                            reader,
375                            value_type.as_ref(),
376                            dict_info.offset,
377                            dict_info.length,
378                            ..,
379                        )
380                        .await?,
381                    );
382                }
383                _ => {
384                    return Err(Error::schema(format!(
385                        "Does not support {} as dictionary value type",
386                        value_type
387                    )));
388                }
389            }
390        } else {
391            panic!("Should not reach here: dictionary field does not load dictionary info")
392        }
393        Ok(())
394    } else {
395        for child in field.children.as_mut_slice() {
396            load_field_dictionary(child, reader).await?;
397        }
398        Ok(())
399    }
400}
401
402/// Load dictionary value array from manifest files.
403// TODO: pub(crate)
404pub async fn populate_schema_dictionary(schema: &mut Schema, reader: &dyn Reader) -> Result<()> {
405    for field in schema.fields.as_mut_slice() {
406        load_field_dictionary(field, reader).await?;
407    }
408    Ok(())
409}
410
411#[cfg(test)]
412mod tests {
413    use std::collections::HashMap;
414
415    use arrow_schema::DataType;
416    use arrow_schema::Field as ArrowField;
417    use arrow_schema::Fields as ArrowFields;
418    use arrow_schema::Schema as ArrowSchema;
419    use lance_core::Error;
420    use lance_core::datatypes::Schema;
421
422    use super::{Fields, FieldsWithMeta};
423    use crate::format::pb;
424
425    fn proto_field(id: i32, parent_id: i32, name: String, logical_type: &str) -> pb::Field {
426        pb::Field {
427            id,
428            parent_id,
429            name,
430            logical_type: logical_type.to_owned(),
431            ..Default::default()
432        }
433    }
434
435    #[test]
436    fn test_schema_set_ids() {
437        let arrow_schema = ArrowSchema::new(vec![
438            ArrowField::new("a", DataType::Int32, false),
439            ArrowField::new(
440                "b",
441                DataType::Struct(ArrowFields::from(vec![
442                    ArrowField::new("f1", DataType::Utf8, true),
443                    ArrowField::new("f2", DataType::Boolean, false),
444                    ArrowField::new("f3", DataType::Float32, false),
445                ])),
446                true,
447            ),
448            ArrowField::new("c", DataType::Float64, false),
449        ]);
450        let schema = Schema::try_from(&arrow_schema).unwrap();
451
452        let protos: Fields = (&schema).into();
453        assert_eq!(
454            protos.0.iter().map(|p| p.id).collect::<Vec<_>>(),
455            (0..6).collect::<Vec<_>>()
456        );
457    }
458
459    #[test]
460    fn test_schema_metadata() {
461        let mut metadata: HashMap<String, String> = HashMap::new();
462        metadata.insert(String::from("k1"), String::from("v1"));
463        metadata.insert(String::from("k2"), String::from("v2"));
464
465        let arrow_schema = ArrowSchema::new_with_metadata(
466            vec![ArrowField::new("a", DataType::Int32, false)],
467            metadata,
468        );
469
470        let expected_schema = Schema::try_from(&arrow_schema).unwrap();
471        let fields_with_meta: FieldsWithMeta = (&expected_schema).into();
472
473        let schema = Schema::try_from(fields_with_meta).unwrap();
474        assert_eq!(expected_schema, schema);
475    }
476
477    #[test]
478    fn test_reconstruct_wide_nested_schema() {
479        const NUM_STRUCTS: usize = 4096;
480
481        let mut proto_fields = Vec::with_capacity(NUM_STRUCTS * 3);
482        for struct_index in 0..NUM_STRUCTS {
483            let parent_id = (struct_index * 3) as i32;
484            proto_fields.push(proto_field(
485                parent_id,
486                -1,
487                format!("struct_{struct_index}"),
488                "struct",
489            ));
490            proto_fields.push(proto_field(
491                parent_id + 1,
492                parent_id,
493                format!("left_{struct_index}"),
494                "int32",
495            ));
496            proto_fields.push(proto_field(
497                parent_id + 2,
498                parent_id,
499                format!("right_{struct_index}"),
500                "int32",
501            ));
502        }
503
504        let fields = Fields(proto_fields);
505        let schema = Schema::try_from(&fields).unwrap();
506        assert_eq!(schema.fields.len(), NUM_STRUCTS);
507        for (struct_index, field) in schema.fields.iter().enumerate() {
508            let parent_id = (struct_index * 3) as i32;
509            assert_eq!(field.id, parent_id);
510            assert_eq!(field.name, format!("struct_{struct_index}"));
511            assert_eq!(field.children.len(), 2);
512            assert_eq!(field.children[0].id, parent_id + 1);
513            assert_eq!(field.children[0].name, format!("left_{struct_index}"));
514            assert_eq!(field.children[1].id, parent_id + 2);
515            assert_eq!(field.children[1].name, format!("right_{struct_index}"));
516        }
517    }
518
519    #[test]
520    fn test_reconstruct_deep_nested_schema() {
521        const DEPTH: usize = 1024;
522
523        let proto_fields = (0..DEPTH)
524            .map(|depth| {
525                proto_field(
526                    depth as i32,
527                    if depth == 0 { -1 } else { depth as i32 - 1 },
528                    format!("level_{depth}"),
529                    if depth + 1 == DEPTH {
530                        "int32"
531                    } else {
532                        "struct"
533                    },
534                )
535            })
536            .collect();
537
538        let fields = Fields(proto_fields);
539        let schema = Schema::try_from(&fields).unwrap();
540        assert_eq!(schema.fields.len(), 1);
541        let mut field = &schema.fields[0];
542        for depth in 0..DEPTH {
543            assert_eq!(field.id, depth as i32);
544            assert_eq!(field.name, format!("level_{depth}"));
545            if depth + 1 == DEPTH {
546                assert!(field.children.is_empty());
547            } else {
548                assert_eq!(field.children.len(), 1);
549                field = &field.children[0];
550            }
551        }
552    }
553
554    #[test]
555    fn test_reconstruct_schema_reports_missing_parent() {
556        let fields = Fields(vec![proto_field(7, 42, "child".to_owned(), "int32")]);
557
558        let error = Schema::try_from(&fields).unwrap_err();
559        assert!(matches!(&error, Error::Schema { .. }));
560        assert!(
561            error.to_string().contains(
562                "Field 'child' (id=7) references parent id 42, which must appear earlier"
563            )
564        );
565    }
566
567    #[test]
568    fn test_reconstruct_schema_preserves_legacy_duplicate_id_match() {
569        let fields = Fields(vec![
570            proto_field(1, -1, "root_a".to_owned(), "struct"),
571            proto_field(2, -1, "root_b".to_owned(), "struct"),
572            proto_field(2, 1, "nested_duplicate".to_owned(), "struct"),
573            proto_field(3, 2, "child".to_owned(), "int32"),
574        ]);
575
576        let schema = Schema::try_from(&fields).unwrap();
577        assert_eq!(schema.fields.len(), 2);
578        assert_eq!(schema.fields[0].name, "root_a");
579        assert_eq!(schema.fields[0].children.len(), 1);
580        assert_eq!(schema.fields[0].children[0].name, "nested_duplicate");
581        assert_eq!(schema.fields[0].children[0].children.len(), 1);
582        assert_eq!(schema.fields[0].children[0].children[0].name, "child");
583        assert_eq!(schema.fields[1].name, "root_b");
584        assert!(schema.fields[1].children.is_empty());
585    }
586
587    #[test]
588    fn test_clustering_key_roundtrip() {
589        let arrow_schema = ArrowSchema::new(vec![
590            ArrowField::new("region", DataType::Utf8, true).with_metadata(
591                vec![(
592                    "lance-schema:unenforced-clustering-key:position".to_owned(),
593                    "1".to_owned(),
594                )]
595                .into_iter()
596                .collect::<HashMap<_, _>>(),
597            ),
598            ArrowField::new("date", DataType::Int32, false).with_metadata(
599                vec![(
600                    "lance-schema:unenforced-clustering-key:position".to_owned(),
601                    "2".to_owned(),
602                )]
603                .into_iter()
604                .collect::<HashMap<_, _>>(),
605            ),
606            ArrowField::new("value", DataType::Float64, true),
607        ]);
608
609        let schema = Schema::try_from(&arrow_schema).unwrap();
610        let ck = schema.unenforced_clustering_key();
611        assert_eq!(ck.len(), 2);
612        assert_eq!(ck[0].name, "region");
613        assert_eq!(ck[1].name, "date");
614
615        // Round-trip through protobuf
616        let fields_with_meta: FieldsWithMeta = (&schema).into();
617        let restored = Schema::try_from(fields_with_meta).unwrap();
618
619        let ck2 = restored.unenforced_clustering_key();
620        assert_eq!(ck2.len(), 2);
621        assert_eq!(ck2[0].name, "region");
622        assert_eq!(ck2[1].name, "date");
623        assert_eq!(ck2[0].unenforced_clustering_key_position, Some(1));
624        assert_eq!(ck2[1].unenforced_clustering_key_position, Some(2));
625
626        // Non-clustering-key field should not have position
627        let value_field = restored.field("value").unwrap();
628        assert!(!value_field.is_unenforced_clustering_key());
629    }
630}
lance_file/datatypes.rs

lance_file/
datatypes.rs
