lance_file/writer/
field_types.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Checks that arrays handed to a file writer have the Arrow types recorded in
5//! the file schema.
6//!
7//! The file schema is what readers use to decode a column, so an array whose
8//! type or extension differs from it is written in a layout the file does not
9//! describe. Callers own any conversion from their input representation (for
10//! example Arrow JSON text to Lance JSONB); this check only rejects arrays that
11//! skipped it, before any page is encoded.
12
13use std::fmt;
14use std::sync::Arc;
15
16use arrow_array::{ArrayRef, RecordBatch};
17use arrow_schema::{DataType, Field as ArrowField, FieldRef, SchemaRef};
18use lance_arrow::BLOB_V2_EXT_NAME;
19use lance_core::{Error, Result, datatypes::Schema};
20
21/// An array whose Arrow type or extension differs from the field the file
22/// schema records for it.
23///
24/// Returned as the source of an [`Error::InvalidInput`] by the file writers.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct FieldTypeMismatch {
27    /// Dot-separated path of the mismatched field, starting at the top-level
28    /// column.
29    pub field_path: String,
30    /// The field the file schema records.
31    pub expected: FieldRef,
32    /// The field the input carries. A column written on its own is a bare
33    /// array, so it carries no extension.
34    pub actual: FieldRef,
35}
36
37impl fmt::Display for FieldTypeMismatch {
38    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
39        write!(
40            f,
41            "field `{}` does not match the file schema: expected {}, got {}",
42            self.field_path,
43            FieldType(&self.expected),
44            FieldType(&self.actual)
45        )
46    }
47}
48
49impl std::error::Error for FieldTypeMismatch {}
50
51struct FieldType<'a>(&'a ArrowField);
52
53impl fmt::Display for FieldType<'_> {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        write!(f, "{}", self.0.data_type())?;
56        match (
57            self.0.extension_type_name(),
58            self.0.extension_type_metadata(),
59        ) {
60            (Some(name), Some(metadata)) => write!(f, " (extension {name}, metadata {metadata})"),
61            (Some(name), None) => write!(f, " (extension {name})"),
62            _ => write!(f, " (no extension)"),
63        }
64    }
65}
66
67/// The Arrow fields a file schema expects, built once per writer.
68pub struct ExpectedTypes {
69    fields: Vec<FieldRef>,
70    /// The last batch schema that passed. Batches of one stream normally share
71    /// a schema, so later batches are accepted without walking it again.
72    accepted: Option<SchemaRef>,
73}
74
75impl ExpectedTypes {
76    pub fn new(schema: &Schema) -> Self {
77        Self {
78            fields: schema
79                .fields
80                .iter()
81                .map(|field| Arc::new(ArrowField::from(field)))
82                .collect(),
83            accepted: None,
84        }
85    }
86
87    /// Check every column of `batch` that the file schema writes.
88    ///
89    /// Columns are matched by name, the same way the encoders select them. A
90    /// missing column is left for the encoder to report.
91    pub fn check_batch(&mut self, batch: &RecordBatch) -> Result<()> {
92        let batch_schema = batch.schema();
93        if self.accepted.as_ref() == Some(&batch_schema) {
94            return Ok(());
95        }
96        for (index, expected) in self.fields.iter().enumerate() {
97            // Batches usually list columns in schema order; only fall back to
98            // a name search when they do not, so wide schemas stay linear.
99            let actual = batch_schema
100                .fields()
101                .get(index)
102                .filter(|actual| actual.name() == expected.name())
103                .or_else(|| {
104                    let (index, _) = batch_schema.column_with_name(expected.name())?;
105                    batch_schema.fields().get(index)
106                });
107            if let Some(actual) = actual {
108                check_field(expected, actual, true).map_err(into_error)?;
109            }
110        }
111        self.accepted = Some(batch_schema);
112        Ok(())
113    }
114
115    /// Check an array written to the top-level column at `index`.
116    ///
117    /// A bare array carries no field metadata, so the column's own extension
118    /// is not compared; nested extensions live in the array's data type and
119    /// are.
120    pub fn check_column(&self, index: usize, array: &ArrayRef) -> Result<()> {
121        let expected = &self.fields[index];
122        let actual = Arc::new(ArrowField::new(
123            expected.name(),
124            array.data_type().clone(),
125            true,
126        ));
127        check_field(expected, &actual, false).map_err(into_error)
128    }
129}
130
131fn into_error(mismatch: FieldTypeMismatch) -> Error {
132    Error::invalid_input_source(Box::new(mismatch))
133}
134
135/// Names and nullability of nested fields are not compared: the encoders
136/// address children by position and verify nulls against the values.
137fn check_field(
138    expected: &FieldRef,
139    actual: &FieldRef,
140    compare_extension: bool,
141) -> std::result::Result<(), FieldTypeMismatch> {
142    let mismatch = || FieldTypeMismatch {
143        field_path: expected.name().clone(),
144        expected: expected.clone(),
145        actual: actual.clone(),
146    };
147    // Blob v2 has no serialized extension parameters. PyArrow serializes an
148    // empty string, while Rust's logical and prepared field builders omit it.
149    let expected_metadata = expected.extension_type_metadata();
150    let actual_metadata = actual.extension_type_metadata();
151    let metadata_matches = expected_metadata == actual_metadata
152        || (expected.extension_type_name() == Some(BLOB_V2_EXT_NAME)
153            && expected_metadata.unwrap_or_default() == actual_metadata.unwrap_or_default());
154    if compare_extension
155        && (expected.extension_type_name() != actual.extension_type_name() || !metadata_matches)
156    {
157        return Err(mismatch());
158    }
159    let children = match (expected.data_type(), actual.data_type()) {
160        // The file schema records the Blob v2 struct users write, but blob
161        // preprocessing hands the encoder a descriptor struct, which the blob
162        // encoder validates itself.
163        (DataType::Struct(_), DataType::Struct(_))
164            if expected.extension_type_name() == Some(BLOB_V2_EXT_NAME) =>
165        {
166            return Ok(());
167        }
168        (DataType::Struct(expected_children), DataType::Struct(actual_children))
169            if expected_children.len() == actual_children.len() =>
170        {
171            expected_children
172                .iter()
173                .zip(actual_children.iter())
174                .collect()
175        }
176        (DataType::List(expected_item), DataType::List(actual_item))
177        | (DataType::LargeList(expected_item), DataType::LargeList(actual_item))
178        | (DataType::Map(expected_item, _), DataType::Map(actual_item, _)) => {
179            vec![(expected_item, actual_item)]
180        }
181        (
182            DataType::FixedSizeList(expected_item, expected_size),
183            DataType::FixedSizeList(actual_item, actual_size),
184        ) if expected_size == actual_size => vec![(expected_item, actual_item)],
185        // The encoders write view arrays in the equivalent offset layout.
186        (DataType::Utf8, DataType::Utf8View) | (DataType::Binary, DataType::BinaryView) => {
187            return Ok(());
188        }
189        (expected_type, actual_type) if expected_type == actual_type => return Ok(()),
190        _ => return Err(mismatch()),
191    };
192    for (expected_child, actual_child) in children {
193        check_field(expected_child, actual_child, true).map_err(|mut mismatch| {
194            mismatch.field_path = format!("{}.{}", expected.name(), mismatch.field_path);
195            mismatch
196        })?;
197    }
198    Ok(())
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use arrow_array::new_empty_array;
205    use arrow_schema::Schema as ArrowSchema;
206    use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY};
207    use lance_core::datatypes::{BLOB_V2_LOGICAL_MINIMAL_FIELDS, BLOB_V2_PREPARED_TYPE};
208    use rstest::rstest;
209    use std::collections::HashMap;
210
211    #[rstest]
212    #[case::python_to_rust(BLOB_V2_EXT_NAME, BLOB_V2_EXT_NAME, Some(""), None, true)]
213    #[case::rust_to_python(BLOB_V2_EXT_NAME, BLOB_V2_EXT_NAME, None, Some(""), true)]
214    #[case::nonempty_expected(BLOB_V2_EXT_NAME, BLOB_V2_EXT_NAME, Some("parameters"), None, false)]
215    #[case::nonempty_actual(BLOB_V2_EXT_NAME, BLOB_V2_EXT_NAME, None, Some("parameters"), false)]
216    #[case::different_name(BLOB_V2_EXT_NAME, "other.extension", None, None, false)]
217    #[case::other_extension("other.extension", "other.extension", Some(""), None, false)]
218    fn blob_extension_metadata(
219        #[case] expected_name: &str,
220        #[case] actual_name: &str,
221        #[case] expected_metadata: Option<&str>,
222        #[case] actual_metadata: Option<&str>,
223        #[case] accepted: bool,
224    ) {
225        let nested_field = |data_type, name: &str, metadata: Option<&str>| {
226            let mut values = HashMap::from([(ARROW_EXT_NAME_KEY.to_string(), name.to_string())]);
227            if let Some(metadata) = metadata {
228                values.insert(ARROW_EXT_META_KEY.to_string(), metadata.to_string());
229            }
230            ArrowField::new(
231                "parent",
232                DataType::Struct(
233                    vec![ArrowField::new("blob", data_type, true).with_metadata(values)].into(),
234                ),
235                true,
236            )
237        };
238        let expected = nested_field(
239            DataType::Struct(BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone()),
240            expected_name,
241            expected_metadata,
242        );
243        let actual_type = if expected_name == BLOB_V2_EXT_NAME {
244            BLOB_V2_PREPARED_TYPE.clone()
245        } else {
246            DataType::Struct(BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone())
247        };
248        let actual = nested_field(actual_type, actual_name, actual_metadata);
249        let schema = Schema::try_from(&ArrowSchema::new(vec![expected])).unwrap();
250        let array = new_empty_array(actual.data_type());
251        let batch = RecordBatch::try_new(
252            Arc::new(ArrowSchema::new(vec![actual])),
253            vec![array.clone()],
254        )
255        .unwrap();
256        let mut expected_types = ExpectedTypes::new(&schema);
257        for result in [
258            expected_types.check_batch(&batch),
259            expected_types.check_column(0, &array),
260        ] {
261            if accepted {
262                result.unwrap();
263            } else {
264                let error = result.unwrap_err();
265                assert!(matches!(error, Error::InvalidInput { .. }));
266                assert!(error.to_string().contains("parent.blob"), "{error}");
267            }
268        }
269    }
270}
lance_file/writer/field_types.rs

lance_file/writer/
field_types.rs
