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}