datafusion_datasource_avro/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18#![doc(
19    html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg",
20    html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg"
21)]
22#![cfg_attr(docsrs, feature(doc_cfg))]
23// Make sure fast / cheap clones on Arc are explicit:
24// https://github.com/apache/datafusion/issues/11143
25#![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))]
26#![cfg_attr(test, allow(clippy::needless_pass_by_value))]
27
28//! An [Avro](https://avro.apache.org/) based [`FileSource`](datafusion_datasource::file::FileSource) implementation and related functionality.
29
30pub mod file_format;
31pub mod source;
32
33use arrow::datatypes::{DataType, Field, Fields, Schema, UnionFields};
34pub use arrow_avro;
35use arrow_avro::reader::ReaderBuilder;
36pub use file_format::*;
37use std::io::{BufReader, Read};
38use std::sync::Arc;
39
40/// Read Avro schema given a reader
41pub fn read_avro_schema_from_reader<R: Read>(
42    reader: &mut R,
43) -> datafusion_common::Result<Schema> {
44    let avro_reader = ReaderBuilder::new().build(BufReader::new(reader))?;
45    // Avro readers perform strict schema resolution rules (e.g. record identity checks)
46    // that are stricter than DataFusion's table schema handling needs for inferred schemas.
47    // Drop metadata from inferred schemas so runtime batches and inferred table schemas
48    // compare consistently without requiring strict Avro metadata identity.
49    Ok(strip_metadata_from_schema(avro_reader.schema().as_ref()))
50}
51
52fn strip_metadata_from_schema(schema: &Schema) -> Schema {
53    let fields = schema
54        .fields
55        .into_iter()
56        .map(|f| Arc::new(strip_metadata_from_field(f.as_ref())))
57        .collect::<Fields>();
58    // Intentionally drop schema-level metadata
59    Schema::new(fields)
60}
61
62fn strip_metadata_from_field(field: &Field) -> Field {
63    // Intentionally drop field-level metadata
64    Field::new(
65        field.name(),
66        strip_metadata_from_data_type(field.data_type()),
67        field.is_nullable(),
68    )
69}
70
71fn strip_metadata_from_data_type(data_type: &DataType) -> DataType {
72    match data_type {
73        DataType::Struct(fields) => DataType::Struct(
74            fields
75                .iter()
76                .map(|f| Arc::new(strip_metadata_from_field(f.as_ref())))
77                .collect(),
78        ),
79        DataType::List(field) => {
80            DataType::List(Arc::new(strip_metadata_from_field(field.as_ref())))
81        }
82        DataType::LargeList(field) => {
83            DataType::LargeList(Arc::new(strip_metadata_from_field(field.as_ref())))
84        }
85        DataType::FixedSizeList(field, size) => DataType::FixedSizeList(
86            Arc::new(strip_metadata_from_field(field.as_ref())),
87            *size,
88        ),
89        DataType::Map(field, sorted) => {
90            DataType::Map(Arc::new(strip_metadata_from_field(field.as_ref())), *sorted)
91        }
92        DataType::Union(fields, mode) => {
93            let (type_ids, children): (Vec<_>, Vec<_>) = fields
94                .iter()
95                .map(|(type_id, field)| {
96                    (type_id, Arc::new(strip_metadata_from_field(field.as_ref())))
97                })
98                .unzip();
99
100            DataType::Union(
101                UnionFields::try_new(type_ids, children)
102                    .expect("existing union fields should remain valid"),
103                *mode,
104            )
105        }
106        _ => data_type.clone(),
107    }
108}
109
110#[cfg(test)]
111mod test {
112    use super::*;
113    use arrow::datatypes::{DataType, Field, TimeUnit};
114    use datafusion_common::Result as DFResult;
115    use datafusion_common::test_util::arrow_test_data;
116    use std::fs::File;
117
118    fn avro_test_file(name: &str) -> String {
119        format!("{}/avro/{name}", arrow_test_data())
120    }
121
122    #[test]
123    fn test_read_avro_schema_from_reader() -> DFResult<()> {
124        let path = avro_test_file("alltypes_dictionary.avro");
125        let mut file = File::open(&path)?;
126        let file_schema = read_avro_schema_from_reader(&mut file)?;
127
128        let expected_fields = vec![
129            Field::new("id", DataType::Int32, true),
130            Field::new("bool_col", DataType::Boolean, true),
131            Field::new("tinyint_col", DataType::Int32, true),
132            Field::new("smallint_col", DataType::Int32, true),
133            Field::new("int_col", DataType::Int32, true),
134            Field::new("bigint_col", DataType::Int64, true),
135            Field::new("float_col", DataType::Float32, true),
136            Field::new("double_col", DataType::Float64, true),
137            Field::new("date_string_col", DataType::Binary, true),
138            Field::new("string_col", DataType::Binary, true),
139            Field::new(
140                "timestamp_col",
141                DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
142                true,
143            ),
144        ];
145
146        assert_eq!(file_schema.fields.len(), expected_fields.len());
147        for (i, field) in file_schema.fields.iter().enumerate() {
148            assert_eq!(field.as_ref(), &expected_fields[i]);
149        }
150
151        Ok(())
152    }
153}
datafusion_datasource_avro/mod.rs

datafusion_datasource_avro/
mod.rs
