datafusion_datasource_csv/
source.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//! Execution plan for reading CSV files
19
20use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
21use datafusion_physical_plan::projection::ProjectionExprs;
22use std::fmt;
23use std::io::{Read, Seek, SeekFrom};
24use std::sync::Arc;
25use std::task::Poll;
26
27use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream};
28use datafusion_datasource::file_compression_type::FileCompressionType;
29use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
30use datafusion_datasource::{
31    FileRange, ListingTableUrl, PartitionedFile, RangeCalculation, TableSchema,
32    as_file_source, calculate_range,
33};
34
35use arrow::csv;
36use datafusion_common::config::CsvOptions;
37use datafusion_common::{DataFusionError, Result};
38use datafusion_common_runtime::JoinSet;
39use datafusion_datasource::file::FileSource;
40use datafusion_datasource::file_scan_config::FileScanConfig;
41use datafusion_execution::TaskContext;
42use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
43use datafusion_physical_plan::{
44    DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
45};
46
47use crate::file_format::CsvDecoder;
48use futures::{StreamExt, TryStreamExt};
49use object_store::buffered::BufWriter;
50use object_store::{GetOptions, GetResultPayload, ObjectStore};
51use tokio::io::AsyncWriteExt;
52
53/// A Config for [`CsvOpener`]
54///
55/// # Example: create a `DataSourceExec` for CSV
56/// ```
57/// # use std::sync::Arc;
58/// # use arrow::datatypes::Schema;
59/// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
60/// # use datafusion_datasource::PartitionedFile;
61/// # use datafusion_datasource_csv::source::CsvSource;
62/// # use datafusion_execution::object_store::ObjectStoreUrl;
63/// # use datafusion_datasource::source::DataSourceExec;
64/// # use datafusion_common::config::CsvOptions;
65///
66/// # let object_store_url = ObjectStoreUrl::local_filesystem();
67/// # let file_schema = Arc::new(Schema::empty());
68///
69/// let options = CsvOptions {
70///     has_header: Some(true),
71///     delimiter: b',',
72///     quote: b'"',
73///     newlines_in_values: Some(true), // The file contains newlines in values
74///     ..Default::default()
75/// };
76/// let source = Arc::new(CsvSource::new(file_schema.clone())
77///     .with_csv_options(options)
78///     .with_terminator(Some(b'#'))
79/// );
80/// // Create a DataSourceExec for reading the first 100MB of `file1.csv`
81/// let config = FileScanConfigBuilder::new(object_store_url, source)
82///     .with_file(PartitionedFile::new("file1.csv", 100*1024*1024))
83///     .build();
84/// let exec = (DataSourceExec::from_data_source(config));
85/// ```
86#[derive(Debug, Clone)]
87pub struct CsvSource {
88    options: CsvOptions,
89    batch_size: Option<usize>,
90    table_schema: TableSchema,
91    projection: SplitProjection,
92    metrics: ExecutionPlanMetricsSet,
93}
94
95impl CsvSource {
96    /// Returns a [`CsvSource`]
97    pub fn new(table_schema: impl Into<TableSchema>) -> Self {
98        let table_schema = table_schema.into();
99        Self {
100            options: CsvOptions::default(),
101            projection: SplitProjection::unprojected(&table_schema),
102            table_schema,
103            batch_size: None,
104            metrics: ExecutionPlanMetricsSet::new(),
105        }
106    }
107
108    /// Sets the CSV options
109    pub fn with_csv_options(mut self, options: CsvOptions) -> Self {
110        self.options = options;
111        self
112    }
113
114    /// true if the first line of each file is a header
115    pub fn has_header(&self) -> bool {
116        self.options.has_header.unwrap_or(true)
117    }
118
119    // true if rows length support truncate
120    pub fn truncate_rows(&self) -> bool {
121        self.options.truncated_rows.unwrap_or(false)
122    }
123    /// A column delimiter
124    pub fn delimiter(&self) -> u8 {
125        self.options.delimiter
126    }
127
128    /// The quote character
129    pub fn quote(&self) -> u8 {
130        self.options.quote
131    }
132
133    /// The line terminator
134    pub fn terminator(&self) -> Option<u8> {
135        self.options.terminator
136    }
137
138    /// Lines beginning with this byte are ignored.
139    pub fn comment(&self) -> Option<u8> {
140        self.options.comment
141    }
142
143    /// The escape character
144    pub fn escape(&self) -> Option<u8> {
145        self.options.escape
146    }
147
148    /// Initialize a CsvSource with escape
149    pub fn with_escape(&self, escape: Option<u8>) -> Self {
150        let mut conf = self.clone();
151        conf.options.escape = escape;
152        conf
153    }
154
155    /// Initialize a CsvSource with terminator
156    pub fn with_terminator(&self, terminator: Option<u8>) -> Self {
157        let mut conf = self.clone();
158        conf.options.terminator = terminator;
159        conf
160    }
161
162    /// Initialize a CsvSource with comment
163    pub fn with_comment(&self, comment: Option<u8>) -> Self {
164        let mut conf = self.clone();
165        conf.options.comment = comment;
166        conf
167    }
168
169    /// Whether to support truncate rows when read csv file
170    pub fn with_truncate_rows(&self, truncate_rows: bool) -> Self {
171        let mut conf = self.clone();
172        conf.options.truncated_rows = Some(truncate_rows);
173        conf
174    }
175
176    /// Whether values may contain newline characters
177    pub fn newlines_in_values(&self) -> bool {
178        self.options.newlines_in_values.unwrap_or(false)
179    }
180}
181
182impl CsvSource {
183    fn open<R: Read>(&self, reader: R) -> Result<csv::Reader<R>> {
184        Ok(self.builder().build(reader)?)
185    }
186
187    fn builder(&self) -> csv::ReaderBuilder {
188        let mut builder =
189            csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema()))
190                .with_delimiter(self.delimiter())
191                .with_batch_size(
192                    self.batch_size
193                        .expect("Batch size must be set before initializing builder"),
194                )
195                .with_header(self.has_header())
196                .with_quote(self.quote())
197                .with_truncated_rows(self.truncate_rows());
198        if let Some(terminator) = self.terminator() {
199            builder = builder.with_terminator(terminator);
200        }
201        builder = builder.with_projection(self.projection.file_indices.clone());
202        if let Some(escape) = self.escape() {
203            builder = builder.with_escape(escape)
204        }
205        if let Some(comment) = self.comment() {
206            builder = builder.with_comment(comment);
207        }
208
209        builder
210    }
211}
212
213/// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`]
214pub struct CsvOpener {
215    config: Arc<CsvSource>,
216    file_compression_type: FileCompressionType,
217    object_store: Arc<dyn ObjectStore>,
218    partition_index: usize,
219}
220
221impl CsvOpener {
222    /// Returns a [`CsvOpener`]
223    pub fn new(
224        config: Arc<CsvSource>,
225        file_compression_type: FileCompressionType,
226        object_store: Arc<dyn ObjectStore>,
227    ) -> Self {
228        Self {
229            config,
230            file_compression_type,
231            object_store,
232            partition_index: 0,
233        }
234    }
235}
236
237impl From<CsvSource> for Arc<dyn FileSource> {
238    fn from(source: CsvSource) -> Self {
239        as_file_source(source)
240    }
241}
242
243impl FileSource for CsvSource {
244    fn create_file_opener(
245        &self,
246        object_store: Arc<dyn ObjectStore>,
247        base_config: &FileScanConfig,
248        partition_index: usize,
249    ) -> Result<Arc<dyn FileOpener>> {
250        let mut opener = Arc::new(CsvOpener {
251            config: Arc::new(self.clone()),
252            file_compression_type: base_config.file_compression_type,
253            object_store,
254            partition_index,
255        }) as Arc<dyn FileOpener>;
256        opener = ProjectionOpener::try_new(
257            self.projection.clone(),
258            Arc::clone(&opener),
259            self.table_schema.file_schema(),
260        )?;
261        Ok(opener)
262    }
263
264    fn table_schema(&self) -> &TableSchema {
265        &self.table_schema
266    }
267
268    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
269        let mut conf = self.clone();
270        conf.batch_size = Some(batch_size);
271        Arc::new(conf)
272    }
273
274    fn try_pushdown_projection(
275        &self,
276        projection: &ProjectionExprs,
277    ) -> Result<Option<Arc<dyn FileSource>>> {
278        let mut source = self.clone();
279        let new_projection = self.projection.source.try_merge(projection)?;
280        let split_projection =
281            SplitProjection::new(self.table_schema.file_schema(), &new_projection);
282        source.projection = split_projection;
283        Ok(Some(Arc::new(source)))
284    }
285
286    fn projection(&self) -> Option<&ProjectionExprs> {
287        Some(&self.projection.source)
288    }
289
290    fn metrics(&self) -> &ExecutionPlanMetricsSet {
291        &self.metrics
292    }
293
294    fn file_type(&self) -> &str {
295        "csv"
296    }
297
298    fn supports_repartitioning(&self) -> bool {
299        // Cannot repartition if values may contain newlines, as record
300        // boundaries cannot be determined by byte offset alone
301        !self.options.newlines_in_values.unwrap_or(false)
302    }
303
304    fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
305        match t {
306            DisplayFormatType::Default | DisplayFormatType::Verbose => {
307                write!(f, ", has_header={}", self.has_header())
308            }
309            DisplayFormatType::TreeRender => Ok(()),
310        }
311    }
312}
313
314impl FileOpener for CsvOpener {
315    /// Open a partitioned CSV file.
316    ///
317    /// If `file_meta.range` is `None`, the entire file is opened.
318    /// If `file_meta.range` is `Some(FileRange {start, end})`, this signifies that the partition
319    /// corresponds to the byte range [start, end) within the file.
320    ///
321    /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules
322    /// are applied to determine which lines to read:
323    /// 1. The first line of the partition is the line in which the index of the first character >= `start`.
324    /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides.
325    ///
326    /// Examples:
327    /// Consider the following partitions enclosed by braces `{}`:
328    ///
329    /// {A,1,2,3,4,5,6,7,8,9\n
330    ///  A,1,2,3,4,5,6,7,8,9\n}
331    ///  A,1,2,3,4,5,6,7,8,9\n
332    ///  The lines read would be: [0, 1]
333    ///
334    ///  A,{1,2,3,4,5,6,7,8,9\n
335    ///  A,1,2,3,4,5,6,7,8,9\n
336    ///  A},1,2,3,4,5,6,7,8,9\n
337    ///  The lines read would be: [1, 2]
338    fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
339        // `self.config.has_header` controls whether to skip reading the 1st line header
340        // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle
341        // partition, then don't skip first line
342        let mut csv_has_header = self.config.has_header();
343        if let Some(FileRange { start, .. }) = partitioned_file.range
344            && start != 0
345        {
346            csv_has_header = false;
347        }
348
349        let mut config = (*self.config).clone();
350        config.options.has_header = Some(csv_has_header);
351        config.options.truncated_rows = Some(config.truncate_rows());
352
353        let file_compression_type = self.file_compression_type.to_owned();
354
355        if partitioned_file.range.is_some() {
356            assert!(
357                !file_compression_type.is_compressed(),
358                "Reading compressed .csv in parallel is not supported"
359            );
360        }
361
362        let store = Arc::clone(&self.object_store);
363        let terminator = self.config.terminator();
364
365        let baseline_metrics =
366            BaselineMetrics::new(&self.config.metrics, self.partition_index);
367
368        Ok(Box::pin(async move {
369            // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries)
370
371            let calculated_range =
372                calculate_range(&partitioned_file, &store, terminator).await?;
373
374            let range = match calculated_range {
375                RangeCalculation::Range(None) => None,
376                RangeCalculation::Range(Some(range)) => Some(range.into()),
377                RangeCalculation::TerminateEarly => {
378                    return Ok(
379                        futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed()
380                    );
381                }
382            };
383
384            let options = GetOptions {
385                range,
386                ..Default::default()
387            };
388
389            let result = store
390                .get_opts(&partitioned_file.object_meta.location, options)
391                .await?;
392
393            match result.payload {
394                #[cfg(not(target_arch = "wasm32"))]
395                GetResultPayload::File(mut file, _) => {
396                    let is_whole_file_scanned = partitioned_file.range.is_none();
397                    let decoder = if is_whole_file_scanned {
398                        // Don't seek if no range as breaks FIFO files
399                        file_compression_type.convert_read(file)?
400                    } else {
401                        file.seek(SeekFrom::Start(result.range.start as _))?;
402                        file_compression_type.convert_read(
403                            file.take((result.range.end - result.range.start) as u64),
404                        )?
405                    };
406
407                    let mut reader = config.open(decoder)?;
408
409                    // Use std::iter::from_fn to wrap execution of iterator's next() method.
410                    let iterator = std::iter::from_fn(move || {
411                        let mut timer = baseline_metrics.elapsed_compute().timer();
412                        let result = reader.next();
413                        timer.stop();
414                        result
415                    });
416
417                    Ok(futures::stream::iter(iterator)
418                        .map(|r| r.map_err(Into::into))
419                        .boxed())
420                }
421                GetResultPayload::Stream(s) => {
422                    let decoder = config.builder().build_decoder();
423                    let s = s.map_err(DataFusionError::from);
424                    let input = file_compression_type.convert_stream(s.boxed())?.fuse();
425
426                    let stream = deserialize_stream(
427                        input,
428                        DecoderDeserializer::new(CsvDecoder::new(decoder)),
429                    );
430                    Ok(stream.map_err(Into::into).boxed())
431                }
432            }
433        }))
434    }
435}
436
437pub async fn plan_to_csv(
438    task_ctx: Arc<TaskContext>,
439    plan: Arc<dyn ExecutionPlan>,
440    path: impl AsRef<str>,
441) -> Result<()> {
442    let path = path.as_ref();
443    let parsed = ListingTableUrl::parse(path)?;
444    let object_store_url = parsed.object_store();
445    let store = task_ctx.runtime_env().object_store(&object_store_url)?;
446    let writer_buffer_size = task_ctx
447        .session_config()
448        .options()
449        .execution
450        .objectstore_writer_buffer_size;
451    let mut join_set = JoinSet::new();
452    for i in 0..plan.output_partitioning().partition_count() {
453        let storeref = Arc::clone(&store);
454        let plan: Arc<dyn ExecutionPlan> = Arc::clone(&plan);
455        let filename = format!("{}/part-{i}.csv", parsed.prefix());
456        let file = object_store::path::Path::parse(filename)?;
457
458        let mut stream = plan.execute(i, Arc::clone(&task_ctx))?;
459        join_set.spawn(async move {
460            let mut buf_writer =
461                BufWriter::with_capacity(storeref, file.clone(), writer_buffer_size);
462            let mut buffer = Vec::with_capacity(1024);
463            //only write headers on first iteration
464            let mut write_headers = true;
465            while let Some(batch) = stream.next().await.transpose()? {
466                let mut writer = csv::WriterBuilder::new()
467                    .with_header(write_headers)
468                    .build(buffer);
469                writer.write(&batch)?;
470                buffer = writer.into_inner();
471                buf_writer.write_all(&buffer).await?;
472                buffer.clear();
473                //prevent writing headers more than once
474                write_headers = false;
475            }
476            buf_writer.shutdown().await.map_err(DataFusionError::from)
477        });
478    }
479
480    while let Some(result) = join_set.join_next().await {
481        match result {
482            Ok(res) => res?, // propagate DataFusion error
483            Err(e) => {
484                if e.is_panic() {
485                    std::panic::resume_unwind(e.into_panic());
486                } else {
487                    unreachable!();
488                }
489            }
490        }
491    }
492
493    Ok(())
494}
datafusion_datasource_csv/source.rs

datafusion_datasource_csv/
source.rs
