datafusion_datasource/
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//! Module containing helper methods for the various file formats
19//! See write.rs for write related helper methods
20
21use std::any::Any;
22use std::collections::HashMap;
23use std::fmt;
24use std::sync::Arc;
25
26use crate::file::FileSource;
27use crate::file_compression_type::FileCompressionType;
28use crate::file_scan_config::FileScanConfig;
29use crate::file_sink_config::FileSinkConfig;
30
31use arrow::datatypes::SchemaRef;
32use datafusion_common::file_options::file_type::FileType;
33use datafusion_common::{internal_err, not_impl_err, GetExt, Result, Statistics};
34use datafusion_physical_expr::LexRequirement;
35use datafusion_physical_plan::ExecutionPlan;
36use datafusion_session::Session;
37
38use async_trait::async_trait;
39use object_store::{ObjectMeta, ObjectStore};
40
41/// Default max records to scan to infer the schema
42pub const DEFAULT_SCHEMA_INFER_MAX_RECORD: usize = 1000;
43
44/// This trait abstracts all the file format specific implementations
45/// from the [`TableProvider`]. This helps code re-utilization across
46/// providers that support the same file formats.
47///
48/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
49#[async_trait]
50pub trait FileFormat: Send + Sync + fmt::Debug {
51    /// Returns the table provider as [`Any`](std::any::Any) so that it can be
52    /// downcast to a specific implementation.
53    fn as_any(&self) -> &dyn Any;
54
55    /// Returns the extension for this FileFormat, e.g. "file.csv" -> csv
56    fn get_ext(&self) -> String;
57
58    /// Returns the extension for this FileFormat when compressed, e.g. "file.csv.gz" -> csv
59    fn get_ext_with_compression(
60        &self,
61        _file_compression_type: &FileCompressionType,
62    ) -> Result<String>;
63
64    /// Infer the common schema of the provided objects. The objects will usually
65    /// be analysed up to a given number of records or files (as specified in the
66    /// format config) then give the estimated common schema. This might fail if
67    /// the files have schemas that cannot be merged.
68    async fn infer_schema(
69        &self,
70        state: &dyn Session,
71        store: &Arc<dyn ObjectStore>,
72        objects: &[ObjectMeta],
73    ) -> Result<SchemaRef>;
74
75    /// Infer the statistics for the provided object. The cost and accuracy of the
76    /// estimated statistics might vary greatly between file formats.
77    ///
78    /// `table_schema` is the (combined) schema of the overall table
79    /// and may be a superset of the schema contained in this file.
80    ///
81    /// TODO: should the file source return statistics for only columns referred to in the table schema?
82    async fn infer_stats(
83        &self,
84        state: &dyn Session,
85        store: &Arc<dyn ObjectStore>,
86        table_schema: SchemaRef,
87        object: &ObjectMeta,
88    ) -> Result<Statistics>;
89
90    /// Take a list of files and convert it to the appropriate executor
91    /// according to this file format.
92    async fn create_physical_plan(
93        &self,
94        state: &dyn Session,
95        conf: FileScanConfig,
96    ) -> Result<Arc<dyn ExecutionPlan>>;
97
98    /// Take a list of files and the configuration to convert it to the
99    /// appropriate writer executor according to this file format.
100    async fn create_writer_physical_plan(
101        &self,
102        _input: Arc<dyn ExecutionPlan>,
103        _state: &dyn Session,
104        _conf: FileSinkConfig,
105        _order_requirements: Option<LexRequirement>,
106    ) -> Result<Arc<dyn ExecutionPlan>> {
107        not_impl_err!("Writer not implemented for this format")
108    }
109
110    /// Return the related FileSource such as `CsvSource`, `JsonSource`, etc.
111    fn file_source(&self) -> Arc<dyn FileSource>;
112}
113
114/// Factory for creating [`FileFormat`] instances based on session and command level options
115///
116/// Users can provide their own `FileFormatFactory` to support arbitrary file formats
117pub trait FileFormatFactory: Sync + Send + GetExt + fmt::Debug {
118    /// Initialize a [FileFormat] and configure based on session and command level options
119    fn create(
120        &self,
121        state: &dyn Session,
122        format_options: &HashMap<String, String>,
123    ) -> Result<Arc<dyn FileFormat>>;
124
125    /// Initialize a [FileFormat] with all options set to default values
126    fn default(&self) -> Arc<dyn FileFormat>;
127
128    /// Returns the table source as [`Any`] so that it can be
129    /// downcast to a specific implementation.
130    fn as_any(&self) -> &dyn Any;
131}
132
133/// A container of [FileFormatFactory] which also implements [FileType].
134/// This enables converting a dyn FileFormat to a dyn FileType.
135/// The former trait is a superset of the latter trait, which includes execution time
136/// relevant methods. [FileType] is only used in logical planning and only implements
137/// the subset of methods required during logical planning.
138#[derive(Debug)]
139pub struct DefaultFileType {
140    file_format_factory: Arc<dyn FileFormatFactory>,
141}
142
143impl DefaultFileType {
144    /// Constructs a [DefaultFileType] wrapper from a [FileFormatFactory]
145    pub fn new(file_format_factory: Arc<dyn FileFormatFactory>) -> Self {
146        Self {
147            file_format_factory,
148        }
149    }
150
151    /// get a reference to the inner [FileFormatFactory] struct
152    pub fn as_format_factory(&self) -> &Arc<dyn FileFormatFactory> {
153        &self.file_format_factory
154    }
155}
156
157impl FileType for DefaultFileType {
158    fn as_any(&self) -> &dyn Any {
159        self
160    }
161}
162
163impl fmt::Display for DefaultFileType {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        write!(f, "{:?}", self.file_format_factory)
166    }
167}
168
169impl GetExt for DefaultFileType {
170    fn get_ext(&self) -> String {
171        self.file_format_factory.get_ext()
172    }
173}
174
175/// Converts a [FileFormatFactory] to a [FileType]
176pub fn format_as_file_type(
177    file_format_factory: Arc<dyn FileFormatFactory>,
178) -> Arc<dyn FileType> {
179    Arc::new(DefaultFileType {
180        file_format_factory,
181    })
182}
183
184/// Converts a [FileType] to a [FileFormatFactory].
185/// Returns an error if the [FileType] cannot be
186/// downcasted to a [DefaultFileType].
187pub fn file_type_to_format(
188    file_type: &Arc<dyn FileType>,
189) -> Result<Arc<dyn FileFormatFactory>> {
190    match file_type
191        .as_ref()
192        .as_any()
193        .downcast_ref::<DefaultFileType>()
194    {
195        Some(source) => Ok(Arc::clone(&source.file_format_factory)),
196        _ => internal_err!("FileType was not DefaultFileType"),
197    }
198}