datafusion_datasource_avro/
file_format.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//! Apache Avro [`FileFormat`] abstractions
19use std::collections::HashMap;
20use std::fmt;
21use std::sync::Arc;
22
23use crate::read_avro_schema_from_reader;
24use crate::source::AvroSource;
25
26use arrow::datatypes::Schema;
27use arrow::datatypes::SchemaRef;
28use datafusion_common::DEFAULT_AVRO_EXTENSION;
29use datafusion_common::GetExt;
30use datafusion_common::internal_err;
31use datafusion_common::parsers::CompressionTypeVariant;
32use datafusion_common::{Result, Statistics};
33use datafusion_datasource::file::FileSource;
34use datafusion_datasource::file_compression_type::FileCompressionType;
35use datafusion_datasource::file_format::{FileFormat, FileFormatFactory};
36use datafusion_datasource::file_scan_config::FileScanConfig;
37use datafusion_datasource::source::DataSourceExec;
38use datafusion_physical_plan::ExecutionPlan;
39use datafusion_session::Session;
40
41use async_trait::async_trait;
42use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt};
43
44#[derive(Default)]
45/// Factory struct used to create [`AvroFormat`]
46pub struct AvroFormatFactory;
47
48impl AvroFormatFactory {
49    /// Creates an instance of [`AvroFormatFactory`]
50    pub fn new() -> Self {
51        Self {}
52    }
53}
54
55impl FileFormatFactory for AvroFormatFactory {
56    fn create(
57        &self,
58        _state: &dyn Session,
59        _format_options: &HashMap<String, String>,
60    ) -> Result<Arc<dyn FileFormat>> {
61        Ok(Arc::new(AvroFormat))
62    }
63
64    fn default(&self) -> Arc<dyn FileFormat> {
65        Arc::new(AvroFormat)
66    }
67}
68
69impl fmt::Debug for AvroFormatFactory {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        f.debug_struct("AvroFormatFactory").finish()
72    }
73}
74
75impl GetExt for AvroFormatFactory {
76    fn get_ext(&self) -> String {
77        // Removes the dot, i.e. ".avro" -> "avro"
78        DEFAULT_AVRO_EXTENSION[1..].to_string()
79    }
80}
81
82/// Avro [`FileFormat`] implementation.
83#[derive(Default, Debug)]
84pub struct AvroFormat;
85
86#[async_trait]
87impl FileFormat for AvroFormat {
88    fn get_ext(&self) -> String {
89        AvroFormatFactory::new().get_ext()
90    }
91
92    fn get_ext_with_compression(
93        &self,
94        file_compression_type: &FileCompressionType,
95    ) -> Result<String> {
96        let ext = self.get_ext();
97        match file_compression_type.get_variant() {
98            CompressionTypeVariant::UNCOMPRESSED => Ok(ext),
99            _ => internal_err!("Avro FileFormat does not support compression."),
100        }
101    }
102
103    fn compression_type(&self) -> Option<FileCompressionType> {
104        None
105    }
106
107    async fn infer_schema(
108        &self,
109        _state: &dyn Session,
110        store: &Arc<dyn ObjectStore>,
111        objects: &[ObjectMeta],
112    ) -> Result<SchemaRef> {
113        let mut schemas = vec![];
114        for object in objects {
115            let r = store.as_ref().get(&object.location).await?;
116            let schema = match r.payload {
117                GetResultPayload::File(mut file, _) => {
118                    read_avro_schema_from_reader(&mut file)?
119                }
120                GetResultPayload::Stream(_) => {
121                    // TODO: Fetching entire file to get schema is potentially wasteful
122                    let data = r.bytes().await?;
123                    read_avro_schema_from_reader(&mut data.as_ref())?
124                }
125            };
126            schemas.push(schema);
127        }
128        let merged_schema = Schema::try_merge(schemas)?;
129        Ok(Arc::new(merged_schema))
130    }
131
132    async fn infer_stats(
133        &self,
134        _state: &dyn Session,
135        _store: &Arc<dyn ObjectStore>,
136        table_schema: SchemaRef,
137        _object: &ObjectMeta,
138    ) -> Result<Statistics> {
139        Ok(Statistics::new_unknown(&table_schema))
140    }
141
142    async fn create_physical_plan(
143        &self,
144        _state: &dyn Session,
145        conf: FileScanConfig,
146    ) -> Result<Arc<dyn ExecutionPlan>> {
147        Ok(DataSourceExec::from_data_source(conf))
148    }
149
150    fn file_source(
151        &self,
152        table_schema: datafusion_datasource::TableSchema,
153    ) -> Arc<dyn FileSource> {
154        Arc::new(AvroSource::new(table_schema))
155    }
156}
datafusion_datasource_avro/file_format.rs

datafusion_datasource_avro/
file_format.rs
