parquet/arrow/
file_type.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Arrow extension type for the Parquet [`FILE`] logical type.
6//!
7//! [`FILE`]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#file
8
9use arrow_schema::extension::ExtensionType;
10use arrow_schema::{ArrowError, DataType};
11
12/// Arrow extension type for the Parquet [`FILE`] logical type.
13///
14/// A `FILE` value is a reference to a range of bytes, stored inline or in an
15/// external file. Its Arrow storage type is a struct whose fields are a subset
16/// of the following, each nullable, identified by name, in any order:
17///
18/// | Field          | Arrow type                                  |
19/// |----------------|---------------------------------------------|
20/// | `uri`          | `Utf8`, `LargeUtf8`, or `Utf8View`          |
21/// | `offset`       | `Int64`                                     |
22/// | `size`         | `Int64`                                     |
23/// | `content_type` | `Utf8`, `LargeUtf8`, or `Utf8View`          |
24/// | `checksum`     | `Utf8`, `LargeUtf8`, or `Utf8View`          |
25/// | `inline`       | `Binary`, `LargeBinary`, or `BinaryView`    |
26///
27/// When reading, a `FILE`-annotated group is returned as a struct with this
28/// extension type. When writing, a struct field with this extension type is
29/// written as a `FILE`-annotated group.
30///
31/// If the file has an embedded Arrow schema, a reader takes each field's
32/// metadata from it rather than from the Parquet annotation. A `FILE`-annotated
33/// group that the embedded schema describes as a plain struct is then read
34/// without this extension type. To get the extension type in that case, read
35/// with [`ArrowReaderOptions::with_skip_arrow_metadata`] set to `true`.
36///
37/// The extension name is `arrow.parquet.file`, which is the intended name for
38/// a canonical Arrow extension type; it is not yet standardized.
39///
40/// [`FILE`]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#file
41/// [`ArrowReaderOptions::with_skip_arrow_metadata`]: crate::arrow::arrow_reader::ArrowReaderOptions::with_skip_arrow_metadata
42#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
43pub struct FileType;
44
45impl ExtensionType for FileType {
46    const NAME: &'static str = "arrow.parquet.file";
47
48    type Metadata = &'static str;
49
50    fn metadata(&self) -> &Self::Metadata {
51        &""
52    }
53
54    fn serialize_metadata(&self) -> Option<String> {
55        Some(String::new())
56    }
57
58    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
59        match metadata {
60            None | Some("") => Ok(""),
61            Some(_) => Err(ArrowError::InvalidArgumentError(
62                "FileType extension type expects no metadata".to_owned(),
63            )),
64        }
65    }
66
67    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
68        let DataType::Struct(fields) = data_type else {
69            return Err(ArrowError::InvalidArgumentError(format!(
70                "FileType only supports a struct, got {data_type}"
71            )));
72        };
73        for (i, field) in fields.iter().enumerate() {
74            let name = field.name();
75            if fields[..i].iter().any(|f| f.name() == name) {
76                return Err(ArrowError::InvalidArgumentError(format!(
77                    "FileType struct has duplicate field '{name}'"
78                )));
79            }
80            let type_ok = match name.as_str() {
81                "uri" | "content_type" | "checksum" => matches!(
82                    field.data_type(),
83                    DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
84                ),
85                "offset" | "size" => field.data_type() == &DataType::Int64,
86                "inline" => matches!(
87                    field.data_type(),
88                    DataType::Binary | DataType::LargeBinary | DataType::BinaryView
89                ),
90                _ => {
91                    return Err(ArrowError::InvalidArgumentError(format!(
92                        "FileType struct has unrecognized field '{name}'. \
93                         Valid fields are: uri, offset, size, content_type, checksum, inline"
94                    )));
95                }
96            };
97            if !type_ok {
98                return Err(ArrowError::InvalidArgumentError(format!(
99                    "FileType field '{name}' has unsupported type {}",
100                    field.data_type()
101                )));
102            }
103            if !field.is_nullable() {
104                return Err(ArrowError::InvalidArgumentError(format!(
105                    "FileType field '{name}' must be nullable"
106                )));
107            }
108        }
109        Ok(())
110    }
111
112    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
113        Self.supports_data_type(data_type)?;
114        Ok(Self)
115    }
116
117    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> {
118        Self.supports_data_type(data_type)
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use std::collections::HashMap;
125    use std::sync::Arc;
126
127    use arrow_array::{
128        Array, ArrayRef, BinaryArray, Int64Array, ListArray, RecordBatch, StringArray, StructArray,
129    };
130    use arrow_buffer::OffsetBuffer;
131    use arrow_schema::{Field, Fields, Schema};
132    use bytes::Bytes;
133
134    use super::*;
135    use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
136    use crate::arrow::{ArrowSchemaConverter, ArrowWriter, PARQUET_FIELD_ID_META_KEY};
137    use crate::basic::LogicalType;
138
139    fn file_fields(first_id: Option<i32>) -> Fields {
140        let field = |name: &str, data_type: DataType, offset: i32| {
141            let field = Field::new(name, data_type, true);
142            match first_id {
143                Some(id) => field.with_metadata(HashMap::from([(
144                    PARQUET_FIELD_ID_META_KEY.to_string(),
145                    (id + offset).to_string(),
146                )])),
147                None => field,
148            }
149        };
150        Fields::from(vec![
151            field("uri", DataType::Utf8, 1),
152            field("offset", DataType::Int64, 2),
153            field("size", DataType::Int64, 3),
154            field("content_type", DataType::Utf8, 4),
155            field("checksum", DataType::Utf8, 5),
156            field("inline", DataType::Binary, 6),
157        ])
158    }
159
160    fn file_field(name: &str, fields: Fields) -> Field {
161        Field::new(name, DataType::Struct(fields), true).with_extension_type(FileType)
162    }
163
164    #[test]
165    fn supports_valid_data_types() {
166        let cases = [
167            file_fields(None),
168            Fields::from(vec![
169                Field::new("size", DataType::Int64, true),
170                Field::new("uri", DataType::Utf8View, true),
171            ]),
172            Fields::from(vec![
173                Field::new("uri", DataType::LargeUtf8, true),
174                Field::new("inline", DataType::BinaryView, true),
175            ]),
176        ];
177        for fields in cases {
178            FileType
179                .supports_data_type(&DataType::Struct(fields))
180                .unwrap();
181        }
182    }
183
184    #[test]
185    fn rejects_invalid_data_types() {
186        let cases = [
187            (DataType::Utf8, "only supports a struct"),
188            (
189                DataType::Struct(Fields::from(vec![Field::new("path", DataType::Utf8, true)])),
190                "unrecognized field 'path'",
191            ),
192            (
193                DataType::Struct(Fields::from(vec![Field::new(
194                    "size",
195                    DataType::Int32,
196                    true,
197                )])),
198                "field 'size' has unsupported type",
199            ),
200            (
201                DataType::Struct(Fields::from(vec![Field::new("uri", DataType::Utf8, false)])),
202                "field 'uri' must be nullable",
203            ),
204            (
205                DataType::Struct(Fields::from(vec![
206                    Field::new("uri", DataType::Utf8, true),
207                    Field::new("uri", DataType::Utf8, true),
208                ])),
209                "duplicate field 'uri'",
210            ),
211        ];
212        for (data_type, expected) in cases {
213            let error = FileType
214                .supports_data_type(&data_type)
215                .unwrap_err()
216                .to_string();
217            assert!(error.contains(expected), "{data_type}: {error}");
218        }
219    }
220
221    #[test]
222    fn schema_round_trip() {
223        let arrow_schema = Schema::new(vec![
224            Field::new("f", DataType::Struct(file_fields(Some(1))), true)
225                .with_metadata(HashMap::from([(
226                    PARQUET_FIELD_ID_META_KEY.to_string(),
227                    "1".to_string(),
228                )]))
229                .with_extension_type(FileType),
230        ]);
231        let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema).unwrap();
232        let group = &parquet_schema.root_schema().get_fields()[0];
233        assert_eq!(
234            group.get_basic_info().logical_type_ref(),
235            Some(&LogicalType::File)
236        );
237        assert_eq!(group.get_basic_info().id(), 1);
238        let ids: Vec<_> = group
239            .get_fields()
240            .iter()
241            .map(|f| f.get_basic_info().id())
242            .collect();
243        assert_eq!(ids, (2..=7).collect::<Vec<_>>());
244
245        let arrow_schema = crate::arrow::parquet_to_arrow_schema(&parquet_schema, None).unwrap();
246        assert_eq!(
247            arrow_schema
248                .field(0)
249                .try_extension_type::<FileType>()
250                .unwrap(),
251            FileType
252        );
253    }
254
255    #[test]
256    fn struct_without_extension_is_not_annotated() {
257        let arrow_schema = Schema::new(vec![Field::new(
258            "f",
259            DataType::Struct(file_fields(None)),
260            true,
261        )]);
262        let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema).unwrap();
263        assert_eq!(
264            parquet_schema.root_schema().get_fields()[0]
265                .get_basic_info()
266                .logical_type_ref(),
267            None
268        );
269    }
270
271    #[test]
272    fn struct_with_unknown_metadata_is_not_annotated() {
273        let field = Field::new("f", DataType::Struct(file_fields(None)), true).with_metadata(
274            HashMap::from([
275                (
276                    "ARROW:extension:name".to_string(),
277                    FileType::NAME.to_string(),
278                ),
279                ("ARROW:extension:metadata".to_string(), "v2".to_string()),
280            ]),
281        );
282        assert!(field.try_extension_type::<FileType>().is_err());
283        let parquet_schema = ArrowSchemaConverter::new()
284            .convert(&Schema::new(vec![field]))
285            .unwrap();
286        assert_eq!(
287            parquet_schema.root_schema().get_fields()[0]
288                .get_basic_info()
289                .logical_type_ref(),
290            None
291        );
292    }
293
294    fn file_values(fields: Fields) -> StructArray {
295        StructArray::new(
296            fields,
297            vec![
298                Arc::new(StringArray::from(vec![Some("s3://b/a.png"), None])) as ArrayRef,
299                Arc::new(Int64Array::from(vec![Some(4), None])),
300                Arc::new(Int64Array::from(vec![Some(10), None])),
301                Arc::new(StringArray::from(vec![Some("image/png"), None])),
302                Arc::new(StringArray::from(vec![None::<&str>, None])),
303                Arc::new(BinaryArray::from(vec![None, Some(b"xyz".as_slice())])),
304            ],
305            None,
306        )
307    }
308
309    fn write_and_read(batch: &RecordBatch) -> (crate::schema::types::TypePtr, RecordBatch) {
310        let mut buffer = Vec::new();
311        let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
312        writer.write(batch).unwrap();
313        writer.close().unwrap();
314        let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap();
315        let root = builder
316            .metadata()
317            .file_metadata()
318            .schema_descr()
319            .root_schema_ptr();
320        let read = builder.build().unwrap().next().unwrap().unwrap();
321        (root, read)
322    }
323
324    #[test]
325    fn writer_round_trip() {
326        let values = file_values(file_fields(None));
327        let schema = Schema::new(vec![file_field("f", values.fields().clone())]);
328        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values.clone())]).unwrap();
329
330        let (root, read) = write_and_read(&batch);
331
332        assert_eq!(
333            root.get_fields()[0].get_basic_info().logical_type_ref(),
334            Some(&LogicalType::File)
335        );
336        assert_eq!(
337            read.schema()
338                .field(0)
339                .try_extension_type::<FileType>()
340                .unwrap(),
341            FileType
342        );
343        assert_eq!(read.column(0).as_ref(), &values as &dyn Array);
344    }
345
346    #[test]
347    fn writer_round_trip_in_list() {
348        let values = file_values(file_fields(None));
349        let element = Arc::new(file_field("element", values.fields().clone()));
350        let list = ListArray::new(
351            element,
352            OffsetBuffer::from_lengths([2]),
353            Arc::new(values),
354            None,
355        );
356        let schema = Schema::new(vec![Field::new("a", list.data_type().clone(), true)]);
357        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list)]).unwrap();
358
359        let (root, read) = write_and_read(&batch);
360
361        let element = &root.get_fields()[0].get_fields()[0].get_fields()[0];
362        assert_eq!(
363            element.get_basic_info().logical_type_ref(),
364            Some(&LogicalType::File)
365        );
366        let DataType::List(element) = read.schema().field(0).data_type().clone() else {
367            panic!("expected a list");
368        };
369        assert_eq!(element.try_extension_type::<FileType>().unwrap(), FileType);
370    }
371}
parquet/arrow/file_type.rs

parquet/arrow/
file_type.rs
